package evolve

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeLLM is a pluggable LLMClient for tests. The handler closure inspects
// the prompt/opts and returns a synthesised response.
type fakeLLM struct {
	mu      sync.Mutex
	calls   int
	prompts []string
	opts    []LLMCallOpts
	handler func(prompt string, opts LLMCallOpts) (string, error)
}

func (f *fakeLLM) Complete(ctx context.Context, prompt string, opts LLMCallOpts) (string, error) {
	f.mu.Lock()
	f.calls++
	f.prompts = append(f.prompts, prompt)
	f.opts = append(f.opts, opts)
	f.mu.Unlock()
	if ctx.Err() != nil {
		return "", ctx.Err()
	}
	if f.handler == nil {
		return "[]", nil
	}
	return f.handler(prompt, opts)
}

func newGen(t *testing.T, f *fakeLLM, opts ...GeneratorOption) *LLMGenerator {
	t.Helper()
	g, err := NewLLMGenerator(f, opts...)
	if err != nil {
		t.Fatalf("NewLLMGenerator: %v", err)
	}
	return g
}

func TestLLMGenerator_HappyPath(t *testing.T) {
	f := &fakeLLM{handler: func(p string, _ LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":{"x":1}},{"id":"b","payload":{"x":2}},{"id":"c","payload":{"x":3}}]`, nil
	}}
	g := newGen(t, f, WithModel("claude-sonnet-4-6"))
	ctx := context.Background()

	cs, err := g.Generate(ctx, "pick parameters", 3)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if len(cs) != 3 {
		t.Fatalf("len: want 3, got %d", len(cs))
	}
	for _, c := range cs {
		if c.ID == "" {
			t.Errorf("empty ID: %+v", c)
		}
		if c.Meta["model"] != "claude-sonnet-4-6" {
			t.Errorf("meta model: got %v", c.Meta["model"])
		}
		if _, ok := c.Meta["generated_at"].(time.Time); !ok {
			t.Errorf("meta generated_at: want time.Time, got %T", c.Meta["generated_at"])
		}
	}
}

func TestLLMGenerator_MarkdownFenceStripped(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return "```json\n[{\"id\":\"a\",\"payload\":1}]\n```", nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), nil, 1)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if len(cs) != 1 || cs[0].ID != "a" {
		t.Errorf("fence strip: got %+v", cs)
	}
}

func TestLLMGenerator_CommentaryAroundJSON(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `Here are 2 candidates for your consideration: [{"id":"a","payload":1},{"id":"b","payload":2}]. Hope that helps!`, nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), "x", 2)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if len(cs) != 2 {
		t.Errorf("commentary around JSON: want 2, got %d", len(cs))
	}
}

func TestLLMGenerator_NestedArrays(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":[[1,2],[3,4]]}]`, nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), "x", 1)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if len(cs) != 1 {
		t.Fatalf("nested arrays: want 1, got %d", len(cs))
	}
}

func TestLLMGenerator_ReturnedFewerThanK(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":1}]`, nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), "x", 5)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	// Contract: Generator returns actual count, no padding.
	if len(cs) != 1 {
		t.Errorf("under-delivery: want 1 (actual), got %d", len(cs))
	}
}

func TestLLMGenerator_ReturnedMoreThanK(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":1},{"id":"b","payload":2},{"id":"c","payload":3}]`, nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), "x", 2)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	// Contract: no truncation.
	if len(cs) != 3 {
		t.Errorf("over-delivery: want 3 (actual), got %d", len(cs))
	}
}

func TestLLMGenerator_InvalidJSON(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return "this is just prose without any JSON", nil
	}}
	g := newGen(t, f)
	_, err := g.Generate(context.Background(), "x", 3)
	if !errors.Is(err, ErrCandidateParseFailed) {
		t.Errorf("invalid JSON: want ErrCandidateParseFailed, got %v", err)
	}
}

func TestLLMGenerator_MalformedArray(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":1`, nil // unclosed array
	}}
	g := newGen(t, f)
	_, err := g.Generate(context.Background(), "x", 1)
	if !errors.Is(err, ErrCandidateParseFailed) {
		t.Errorf("malformed array: want ErrCandidateParseFailed, got %v", err)
	}
}

func TestLLMGenerator_LLMError(t *testing.T) {
	rootErr := errors.New("backend unavailable")
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return "", rootErr
	}}
	g := newGen(t, f)
	_, err := g.Generate(context.Background(), "x", 3)
	if !errors.Is(err, ErrLLMFailed) {
		t.Errorf("LLM error: want ErrLLMFailed wrap, got %v", err)
	}
}

func TestLLMGenerator_InvalidK(t *testing.T) {
	f := &fakeLLM{}
	g := newGen(t, f)
	if _, err := g.Generate(context.Background(), "x", 0); !errors.Is(err, ErrInvalidK) {
		t.Errorf("K=0: want ErrInvalidK, got %v", err)
	}
	if _, err := g.Generate(context.Background(), "x", -5); !errors.Is(err, ErrInvalidK) {
		t.Errorf("K=-5: want ErrInvalidK, got %v", err)
	}
	if f.calls != 0 {
		t.Errorf("LLM must not be called on invalid K, calls=%d", f.calls)
	}
}

func TestLLMGenerator_MetaAutoFields(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"id":"a","payload":1,"meta":{"custom":"tag","model":"fake-hack"}}]`, nil
	}}
	g := newGen(t, f, WithModel("claude-sonnet-4-6"))
	before := time.Now().UTC()
	cs, err := g.Generate(context.Background(), "x", 1, WithTemperature(0.7), WithRoles("skeptic", "optimist"))
	if err != nil {
		t.Fatal(err)
	}
	after := time.Now().UTC()
	m := cs[0].Meta
	if m["temperature"] != 0.7 {
		t.Errorf("meta temperature: got %v", m["temperature"])
	}
	// Auto field must override LLM-supplied "model".
	if m["model"] != "claude-sonnet-4-6" {
		t.Errorf("meta model (auto-wins): got %v", m["model"])
	}
	// Custom LLM-supplied field retained.
	if m["custom"] != "tag" {
		t.Errorf("meta custom: got %v", m["custom"])
	}
	roles, ok := m["roles"].([]string)
	if !ok || len(roles) != 2 || roles[0] != "skeptic" {
		t.Errorf("meta roles: got %v", m["roles"])
	}
	gen, ok := m["generated_at"].(time.Time)
	if !ok {
		t.Fatalf("meta generated_at type: got %T", m["generated_at"])
	}
	if gen.Before(before) || gen.After(after) {
		t.Errorf("meta generated_at %v outside [%v, %v]", gen, before, after)
	}
}

func TestLLMGenerator_RolesInPrompt(t *testing.T) {
	f := &fakeLLM{handler: func(prompt string, _ LLMCallOpts) (string, error) {
		if !strings.Contains(prompt, "skeptic") || !strings.Contains(prompt, "optimist") {
			t.Errorf("roles missing from prompt: %q", prompt)
		}
		return "[]", nil
	}}
	g := newGen(t, f)
	_, _ = g.Generate(context.Background(), "x", 3, WithRoles("skeptic", "optimist"))
}

func TestLLMGenerator_TemperatureAndMaxTokensForwarded(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) { return "[]", nil }}
	g := newGen(t, f, WithMaxTokens(2048), WithModel("m"))
	_, _ = g.Generate(context.Background(), "x", 1, WithTemperature(0.9))
	if f.opts[0].Temperature != 0.9 {
		t.Errorf("temperature forward: got %v", f.opts[0].Temperature)
	}
	if f.opts[0].MaxTokens != 2048 {
		t.Errorf("max tokens forward: got %d", f.opts[0].MaxTokens)
	}
	if f.opts[0].Model != "m" {
		t.Errorf("model forward: got %q", f.opts[0].Model)
	}
}

func TestLLMGenerator_CustomTemplate(t *testing.T) {
	tmpl := "GEN-{{.K}}:{{.Input}}"
	f := &fakeLLM{handler: func(prompt string, _ LLMCallOpts) (string, error) {
		if prompt != "GEN-3:hello" {
			t.Errorf("custom template: got %q", prompt)
		}
		return "[]", nil
	}}
	g := newGen(t, f, WithPromptTemplate(tmpl))
	_, _ = g.Generate(context.Background(), "hello", 3)
}

func TestLLMGenerator_StringInputDirect(t *testing.T) {
	f := &fakeLLM{handler: func(prompt string, _ LLMCallOpts) (string, error) {
		if !strings.Contains(prompt, "raw-text-input") {
			t.Errorf("string input should appear verbatim: %q", prompt)
		}
		return "[]", nil
	}}
	g := newGen(t, f)
	_, _ = g.Generate(context.Background(), "raw-text-input", 1)
}

func TestLLMGenerator_StructInputJSONEncoded(t *testing.T) {
	type task struct {
		Name string
		Qty  int
	}
	f := &fakeLLM{handler: func(prompt string, _ LLMCallOpts) (string, error) {
		if !strings.Contains(prompt, `"Name":"bulk"`) || !strings.Contains(prompt, `"Qty":42`) {
			t.Errorf("struct input JSON-encoded missing: %q", prompt)
		}
		return "[]", nil
	}}
	g := newGen(t, f)
	_, _ = g.Generate(context.Background(), task{Name: "bulk", Qty: 42}, 1)
}

func TestLLMGenerator_CtxCanceled(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) { return "[]", nil }}
	g := newGen(t, f)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := g.Generate(ctx, "x", 1)
	if !errors.Is(err, context.Canceled) {
		t.Errorf("ctx canceled: want context.Canceled, got %v", err)
	}
	if f.calls != 0 {
		t.Errorf("LLM must not be called with canceled ctx, calls=%d", f.calls)
	}
}

func TestLLMGenerator_NilClientRejected(t *testing.T) {
	if _, err := NewLLMGenerator(nil); err == nil {
		t.Errorf("nil client: want error, got nil")
	}
}

func TestLLMGenerator_BadTemplateRejected(t *testing.T) {
	f := &fakeLLM{}
	if _, err := NewLLMGenerator(f, WithPromptTemplate("{{.Bad")); err == nil {
		t.Errorf("bad template: want error, got nil")
	}
}

func TestLLMGenerator_MaxTokensNegative(t *testing.T) {
	f := &fakeLLM{}
	if _, err := NewLLMGenerator(f, WithMaxTokens(-1)); err == nil {
		t.Errorf("negative max tokens: want error, got nil")
	}
}

func TestLLMGenerator_AutoIDWhenMissing(t *testing.T) {
	f := &fakeLLM{handler: func(string, LLMCallOpts) (string, error) {
		return `[{"payload":1},{"id":"","payload":2},{"id":"c","payload":3}]`, nil
	}}
	g := newGen(t, f)
	cs, err := g.Generate(context.Background(), "x", 3)
	if err != nil {
		t.Fatal(err)
	}
	if cs[0].ID != "cand-1" || cs[1].ID != "cand-2" || cs[2].ID != "c" {
		t.Errorf("auto ID fallback: got %q %q %q", cs[0].ID, cs[1].ID, cs[2].ID)
	}
}

func TestExtractJSONArray_Depth(t *testing.T) {
	// String containing "]" should not close the outer array early.
	got := extractJSONArray(`prefix [{"k":"hello]world"}] suffix`)
	if got != `[{"k":"hello]world"}]` {
		t.Errorf("depth-aware extract: got %q", got)
	}
}

func TestExtractJSONArray_Escape(t *testing.T) {
	got := extractJSONArray(`[{"k":"a\"b"}]`)
	if got != `[{"k":"a\"b"}]` {
		t.Errorf("escaped quote extract: got %q", got)
	}
}