package agent

import (
	"context"
	"encoding/json"
	"iter"
	"testing"
	"time"

	"nub/internal/llm"
	"nub/internal/session"
	"nub/internal/tool"
)

// fakeProvider spielt eine vordefinierte Folge von Event-Batches ab, eine
// pro Stream()-Aufruf. Damit lässt sich Multi-Turn-Verhalten ohne Netzwerk
// testen (Teststrategie, Abschnitt 7).
type fakeProvider struct {
	batches [][]llm.Event
	call    int
	caps    llm.Caps
}

func (p *fakeProvider) Name() string   { return "fake" }
func (p *fakeProvider) Caps() llm.Caps { return p.caps }

func (p *fakeProvider) Stream(ctx context.Context, req llm.Request) (iter.Seq2[llm.Event, error], error) {
	if p.call >= len(p.batches) {
		return func(yield func(llm.Event, error) bool) {
			yield(llm.Done{Stop: llm.StopEnd}, nil)
		}, nil
	}
	batch := p.batches[p.call]
	p.call++
	return func(yield func(llm.Event, error) bool) {
		for _, ev := range batch {
			if !yield(ev, nil) {
				return
			}
		}
	}, nil
}

type echoTool struct{}

func (echoTool) Name() string            { return "echo" }
func (echoTool) Description() string     { return "echoes input" }
func (echoTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
func (echoTool) Run(ctx context.Context, input json.RawMessage, env tool.Env) (tool.Result, error) {
	return tool.Result{ForModel: "ok:" + string(input)}, nil
}

func textEvents(s string) []llm.Event {
	return []llm.Event{
		llm.BlockStart{Index: 0, Block: llm.Block{Kind: llm.KindText}},
		llm.BlockDelta{Index: 0, Text: s},
		llm.BlockStop{Index: 0},
		llm.Done{Stop: llm.StopEnd, Usage: llm.Usage{OutputTokens: 1}},
	}
}

func toolCallEvents(id, name, args string) []llm.Event {
	return []llm.Event{
		llm.BlockStart{Index: 0, Block: llm.Block{Kind: llm.KindToolUse, ID: id, Name: name}},
		llm.BlockDelta{Index: 0, PartialJSON: args},
		llm.BlockStop{Index: 0},
		llm.Done{Stop: llm.StopToolUse, Usage: llm.Usage{OutputTokens: 1}},
	}
}

func newTestLoop(provider llm.Provider) (*Loop, *tool.Registry) {
	reg := tool.NewRegistry()
	reg.Register(echoTool{})
	return &Loop{
		Provider: provider,
		Tools:    reg,
		Model:    "test-model",
		Env:      tool.Env{Cwd: ".", RepoRoot: "."},
		Session:  session.New("test-session"),
	}, reg
}

func drain(t *testing.T, out <-chan tool.UIEvent, timeout time.Duration) []tool.UIEvent {
	t.Helper()
	var events []tool.UIEvent
	deadline := time.After(timeout)
	for {
		select {
		case ev, ok := <-out:
			if !ok {
				return events
			}
			events = append(events, ev)
		case <-deadline:
			t.Fatal("timed out waiting for events")
			return nil
		}
	}
}

func TestLoop_SimpleTextTurn(t *testing.T) {
	provider := &fakeProvider{batches: [][]llm.Event{textEvents("hello")}}
	loop, _ := newTestLoop(provider)

	in := make(chan Input, 1)
	in <- Input{Text: "hi"}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	out := loop.Run(ctx, in)
	close(in)
	events := drain(t, out, 3*time.Second)

	var gotText string
	var gotDone bool
	for _, ev := range events {
		switch e := ev.(type) {
		case tool.TextDelta:
			gotText += e.Text
		case tool.TurnDone:
			gotDone = true
			if e.Stop != llm.StopEnd {
				t.Errorf("stop = %q, want end_turn", e.Stop)
			}
		}
	}
	if gotText != "hello" {
		t.Errorf("text = %q", gotText)
	}
	if !gotDone {
		t.Error("expected TurnDone event")
	}
}

func TestLoop_ToolCallRoundTrip(t *testing.T) {
	provider := &fakeProvider{batches: [][]llm.Event{
		toolCallEvents("call_1", "echo", `{"x":1}`),
		textEvents("done"),
	}}
	loop, _ := newTestLoop(provider)

	in := make(chan Input, 1)
	in <- Input{Text: "run echo"}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	out := loop.Run(ctx, in)
	close(in)
	events := drain(t, out, 3*time.Second)

	var gotOutput *tool.ToolCallOutput
	for _, ev := range events {
		if o, ok := ev.(tool.ToolCallOutput); ok {
			cp := o
			gotOutput = &cp
		}
	}
	if gotOutput == nil {
		t.Fatal("expected a ToolCallOutput event")
	}
	if gotOutput.Result.ForModel != `ok:{"x":1}` {
		t.Errorf("tool result = %q", gotOutput.Result.ForModel)
	}
	if provider.call != 2 {
		t.Errorf("expected 2 provider calls (tool_use + follow-up), got %d", provider.call)
	}
}

func TestLoop_UnknownToolProducesErrorResult(t *testing.T) {
	provider := &fakeProvider{batches: [][]llm.Event{
		toolCallEvents("call_1", "does_not_exist", `{}`),
		textEvents("done"),
	}}
	loop, _ := newTestLoop(provider)

	in := make(chan Input, 1)
	in <- Input{Text: "go"}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	out := loop.Run(ctx, in)
	close(in)
	events := drain(t, out, 3*time.Second)

	found := false
	for _, ev := range events {
		if o, ok := ev.(tool.ToolCallOutput); ok && o.Name == "does_not_exist" {
			found = true
			if !o.Result.IsError {
				t.Error("expected IsError for unknown tool")
			}
		}
	}
	if !found {
		t.Error("expected ToolCallOutput for unknown tool")
	}
}

func TestLoop_StreamErrorAbortsWithoutIncompleteMessage(t *testing.T) {
	provider := &fakeProvider{}
	loop, _ := newTestLoop(provider)
	loop.Provider = errorProvider{}

	in := make(chan Input, 1)
	in <- Input{Text: "hi"}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	out := loop.Run(ctx, in)
	close(in)
	events := drain(t, out, 3*time.Second)

	sawError := false
	for _, ev := range events {
		if _, ok := ev.(tool.ErrorEvent); ok {
			sawError = true
		}
	}
	if !sawError {
		t.Error("expected ErrorEvent")
	}
}

type errorProvider struct{}

func (errorProvider) Name() string   { return "error" }
func (errorProvider) Caps() llm.Caps { return llm.Caps{} }
func (errorProvider) Stream(ctx context.Context, req llm.Request) (iter.Seq2[llm.Event, error], error) {
	return func(yield func(llm.Event, error) bool) {
		yield(nil, context.DeadlineExceeded)
	}, nil
}