package agent

import (
	"bytes"
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

// newInterruptibleAgent wires an agent whose "wait" tool blocks until the
// context is cancelled, which is what pressing esc does to a turn.
func newInterruptibleAgent(t *testing.T, f *pausableAPI, root string, started chan<- struct{}) *Agent {
	t.Helper()
	log := metrics.NewWriter(&bytes.Buffer{}, "test")
	state := NewState()
	reg := tools.NewRegistry(log).Register(
		tools.NewRead(root, state),
		tools.New("wait", "Blocks until the turn is cancelled", func(ctx context.Context, _ struct{}) (string, error) {
			started <- struct{}{}
			<-ctx.Done()
			return "", ctx.Err()
		}),
	)
	client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"))
	return New(client, reg, state, log, Options{})
}

// unanswered returns the tool_use ids in msgs that no tool_result answers. The
// API rejects the whole request over even one of them, so this is the only
// property that matters about an interrupted conversation.
func unanswered(msgs []llm.Message) []string {
	answered := map[string]bool{}
	for _, m := range msgs {
		for _, b := range m.Content {
			if b.Type == llm.BlockToolResult {
				answered[b.ToolUseID] = true
			}
		}
	}
	var open []string
	for _, m := range msgs {
		for _, b := range m.Content {
			if b.Type == llm.BlockToolUse && !answered[b.ID] {
				open = append(open, b.ID)
			}
		}
	}
	return open
}

// Interrupting a turn while a tool is in flight must still leave a conversation
// the API will accept: the assistant message asking for the call is already in
// the history, and a tool_use with no tool_result is a 400 on every request
// that follows — the next prompt, or a --resume weeks later.
func TestInterruptedTurnAnswersEveryToolCall(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("A\n"), 0o644)

	f := newPausableAPI(t,
		msg("tool_use",
			toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`),
			toolUseBlock("tu_2", "wait", `{}`)),
		msg("end_turn", textBlock("ok")),
	)
	started := make(chan struct{}, 1)
	a := newInterruptibleAgent(t, f, root, started)

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		_, err := a.Run(ctx, "go")
		done <- err
	}()

	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("the blocking tool never started")
	}
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("the cancelled run never returned")
	}

	h := a.State().History()
	if open := unanswered(h); len(open) != 0 {
		t.Fatalf("after an interrupt these tool_use ids have no tool_result: %v", open)
	}

	last := h[len(h)-1]
	if last.Role != llm.RoleUser || len(last.Content) != 2 {
		t.Fatalf("last message = %+v, want one user message holding both results", last)
	}
	// The call that did finish keeps what it returned: an interrupt throws
	// away the turn, not the work the turn already did.
	if got := last.Content[0]; got.ToolUseID != "tu_1" || got.IsError || !strings.Contains(got.Content, "A") {
		t.Errorf("result for the finished read = %+v, want its contents", got)
	}
	if got := last.Content[1]; got.ToolUseID != "tu_2" || !got.IsError || !strings.Contains(got.Content, "interrupted") {
		t.Errorf("result for the interrupted call = %+v, want an error saying it was interrupted", got)
	}
}

// The repaired conversation is not merely well formed in memory: the next
// prompt in the same session has to reach the API.
func TestTheTurnAfterAnInterruptIsSent(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("A\n"), 0o644)

	f := newPausableAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "wait", `{}`)),
		msg("end_turn", textBlock("ok")),
	)
	started := make(chan struct{}, 1)
	a := newInterruptibleAgent(t, f, root, started)

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		a.Run(ctx, "go")
		close(done)
	}()
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("the blocking tool never started")
	}
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("the cancelled run never returned")
	}

	if _, err := a.Run(context.Background(), "continue pls"); err != nil {
		t.Fatalf("the prompt after an interrupt failed: %v", err)
	}
	if open := unanswered(f.lastRequest(t).Messages); len(open) != 0 {
		t.Errorf("the request sent after an interrupt left these tool_use ids unanswered: %v", open)
	}
}

// A session saved while a turn was in flight is on disk with the dangling
// tool_use in it. Reopening it must repair the history rather than hand the
// user a 400 on their first word.
func TestRestoreClosesADanglingToolUse(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("ok")))
	a := newInterruptibleAgent(t, f, t.TempDir(), make(chan struct{}, 1))

	a.Restore(Session{Messages: []llm.Message{
		{Role: llm.RoleUser, Content: []llm.ContentBlock{{Type: llm.BlockText, Text: "go"}}},
		{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
			{Type: llm.BlockToolUse, ID: "tu_1", Name: "wait", Input: []byte(`{}`)},
		}},
	}})

	h := a.State().History()
	if open := unanswered(h); len(open) != 0 {
		t.Fatalf("restoring left these tool_use ids unanswered: %v", open)
	}
	if last := h[len(h)-1]; last.Role != llm.RoleUser || last.Content[0].ToolUseID != "tu_1" {
		t.Errorf("last message = %+v, want a tool_result for tu_1", last)
	}
}

// A hosted web search is answered by the API itself, in the same assistant
// message. Closing it out locally would send a tool_result for an id the model
// never asked us to run — an error of its own.
func TestRestoreLeavesHostedToolCallsAlone(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("ok")))
	a := newInterruptibleAgent(t, f, t.TempDir(), make(chan struct{}, 1))

	restored := []llm.Message{
		{Role: llm.RoleUser, Content: []llm.ContentBlock{{Type: llm.BlockText, Text: "search"}}},
		{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
			{Type: llm.BlockServerToolUse, ID: "srvtoolu_1", Name: "web_search", Input: []byte(`{}`)},
		}},
	}
	a.Restore(Session{Messages: restored})

	if got := a.State().History(); len(got) != len(restored) {
		t.Errorf("history has %d messages, want the %d restored unchanged", len(got), len(restored))
	}
}