package tools

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"

	"github.com/johannSo/codey/internal/metrics"
)

type echoArgs struct {
	Text  string `json:"text" jsonschema_description:"Text to echo back"`
	Times int    `json:"times,omitempty" jsonschema_description:"How many times to repeat it"`
}

func echoTool(fail bool) Tool {
	return New("echo", "Echo text back", func(_ context.Context, in echoArgs) (string, error) {
		if fail {
			return "", errors.New("echo refused")
		}
		n := in.Times
		if n == 0 {
			n = 1
		}
		return strings.Repeat(in.Text, n), nil
	})
}

func TestSchemaForReflectsStructTags(t *testing.T) {
	var got struct {
		Type       string   `json:"type"`
		Required   []string `json:"required"`
		Properties map[string]struct {
			Type        string `json:"type"`
			Description string `json:"description"`
		} `json:"properties"`
		AdditionalProperties *bool           `json:"additionalProperties"`
		Schema               json.RawMessage `json:"$schema"`
	}
	if err := json.Unmarshal(SchemaFor[echoArgs](), &got); err != nil {
		t.Fatalf("schema is not valid JSON: %v", err)
	}
	if got.Type != "object" {
		t.Errorf("type = %q, want object", got.Type)
	}
	if len(got.Required) != 1 || got.Required[0] != "text" {
		t.Errorf("required = %v, want [text] (omitempty fields must be optional)", got.Required)
	}
	if got.Properties["times"].Type != "integer" {
		t.Errorf("times type = %q, want integer", got.Properties["times"].Type)
	}
	if got.Properties["text"].Description != "Text to echo back" {
		t.Errorf("description not carried over: %q", got.Properties["text"].Description)
	}
	if got.Schema != nil {
		t.Errorf("$schema should be stripped, got %s", got.Schema)
	}
}

// jsonschema:"description=..." is split on commas, so half of any sentence
// with one in it never reached the model. jsonschema_description takes the tag
// whole, and every tool here uses it for that reason; this is the test that
// notices if one drifts back.
func TestSchemaKeepsDescriptionsPastAComma(t *testing.T) {
	type args struct {
		Path string `json:"path" jsonschema_description:"Path to the file, absolute or relative to the working root"`
	}
	var got struct {
		Properties map[string]struct {
			Description string `json:"description"`
		} `json:"properties"`
	}
	if err := json.Unmarshal(SchemaFor[args](), &got); err != nil {
		t.Fatalf("schema is not valid JSON: %v", err)
	}
	if want := "Path to the file, absolute or relative to the working root"; got.Properties["path"].Description != want {
		t.Errorf("description = %q, want the whole of %q", got.Properties["path"].Description, want)
	}
}

func TestSchemaForNoFields(t *testing.T) {
	type empty struct{}
	var got map[string]any
	if err := json.Unmarshal(SchemaFor[empty](), &got); err != nil {
		t.Fatalf("invalid JSON: %v", err)
	}
	if got["type"] != "object" {
		t.Errorf("type = %v, want object", got["type"])
	}
	if _, ok := got["properties"]; !ok {
		t.Error("properties must be present even when empty")
	}
}

func TestTypedExecute(t *testing.T) {
	tests := []struct {
		name    string
		raw     string
		want    string
		wantErr string
	}{
		{name: "decodes typed args", raw: `{"text":"ab","times":3}`, want: "ababab"},
		{name: "omitted optional field", raw: `{"text":"x"}`, want: "x"},
		{name: "empty input treated as {}", raw: ``, want: ""},
		{name: "malformed JSON", raw: `{"text":`, wantErr: "invalid arguments"},
		{name: "wrong field type", raw: `{"text":"x","times":"three"}`, wantErr: "invalid arguments"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got, err := echoTool(false).Execute(context.Background(), json.RawMessage(tc.raw))
			switch {
			case tc.wantErr != "" && err == nil:
				t.Fatalf("want error containing %q, got none (result %q)", tc.wantErr, got)
			case tc.wantErr != "":
				if !strings.Contains(err.Error(), tc.wantErr) {
					t.Fatalf("error = %v, want it to contain %q", err, tc.wantErr)
				}
			case err != nil:
				t.Fatalf("unexpected error: %v", err)
			case got != tc.want:
				t.Fatalf("result = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestRegistryDispatch(t *testing.T) {
	r := NewRegistry(nil).Register(echoTool(false))

	got, err := r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
	if err != nil || got.Text != "hi" || len(got.Blocks) != 0 {
		t.Fatalf("Execute = (%+v, %v), want text \"hi\" and no blocks", got, err)
	}

	if _, err := r.Execute(context.Background(), "nope", nil); err == nil {
		t.Error("unknown tool must return an error the model can read")
	} else if !strings.Contains(err.Error(), "nope") {
		t.Errorf("error should name the tool, got %v", err)
	}
}

func TestRegistryToolErrorIsReturnedNotSwallowed(t *testing.T) {
	r := NewRegistry(nil).Register(echoTool(true))
	out, err := r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
	if err == nil {
		t.Fatal("want the tool's error to surface")
	}
	if out.Text != "" {
		t.Errorf("output = %q, want empty on failure", out.Text)
	}
	if err.Error() != "echo refused" {
		t.Errorf("error = %q, want it passed through verbatim", err)
	}
}

func TestRegistryListIsStable(t *testing.T) {
	r := NewRegistry(nil).Register(
		New("zulu", "", func(context.Context, echoArgs) (string, error) { return "", nil }),
		New("alpha", "", func(context.Context, echoArgs) (string, error) { return "", nil }),
	)
	var names []string
	for _, tl := range r.List() {
		names = append(names, tl.Name())
	}
	if len(names) != 2 || names[0] != "alpha" || names[1] != "zulu" {
		t.Errorf("List = %v, want sorted [alpha zulu] for a byte-stable prompt prefix", names)
	}
}

func TestRegistryDuplicatePanics(t *testing.T) {
	defer func() {
		if recover() == nil {
			t.Error("registering the same name twice must panic at wiring time")
		}
	}()
	NewRegistry(nil).Register(echoTool(false), echoTool(false))
}

func TestExecuteRespectsContext(t *testing.T) {
	blocked := New("block", "", func(ctx context.Context, _ echoArgs) (string, error) {
		<-ctx.Done()
		return "", ctx.Err()
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := NewRegistry(nil).Register(blocked).Execute(ctx, "block", nil); !errors.Is(err, context.Canceled) {
		t.Errorf("error = %v, want context.Canceled", err)
	}
}

func TestRegistryLogsEveryDispatch(t *testing.T) {
	var buf bytes.Buffer
	r := NewRegistry(metrics.NewWriter(&buf, "sess")).Register(echoTool(false))
	r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
	r.Execute(context.Background(), "missing", nil)

	lines := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
	if len(lines) != 2 {
		t.Fatalf("got %d log lines, want one per dispatch:\n%s", len(lines), buf.String())
	}
	var ok struct {
		Tool string `json:"tool"`
		OK   bool   `json:"ok"`
		Out  int    `json:"output_bytes"`
	}
	json.Unmarshal([]byte(lines[0]), &ok)
	if ok.Tool != "echo" || !ok.OK || ok.Out != 2 {
		t.Errorf("success event = %+v, want echo/ok/2 bytes", ok)
	}
	if !strings.Contains(lines[1], `"ok":false`) {
		t.Errorf("unknown-tool dispatch should log a failure, got %s", lines[1])
	}
}

func TestExcludeKeepsToolsOutBeforeAndAfter(t *testing.T) {
	named := func(name string) Tool {
		return New(name, "x", func(context.Context, struct{}) (string, error) { return "", nil })
	}
	r := NewRegistry(nil).Register(named("bash"), named("read_file"))
	r.Exclude(func(n string) bool { return n == "bash" || n == "web_fetch" })
	r.Exclude(func(n string) bool { return strings.HasPrefix(n, "mcp__db") })
	r.Register(named("web_fetch"), named("mcp__db__query"), named("grep"))

	var got []string
	for _, tl := range r.List() {
		got = append(got, tl.Name())
	}
	if strings.Join(got, ",") != "grep,read_file" {
		t.Errorf("tools = %v, want only grep and read_file", got)
	}
	if _, err := r.Execute(context.Background(), "bash", nil); err == nil {
		t.Error("an excluded tool cannot be executed")
	}
	// Registering an excluded name again is not a duplicate: it was never in.
	r.Register(named("bash"))
}