package agent

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

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	"github.com/VictorVargas/rony-llm-agent/pkg/tools"
)

// scriptedLLM returns one canned response per call, in order.
type scriptedLLM struct {
	responses []llm.CompletionResponse
	calls     int
	// lastMessages records the request messages of the most recent call, so
	// tests can assert the corrective nudge was actually sent.
	lastMessages []llm.Message
}

func (s *scriptedLLM) Generate(_ context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
	s.lastMessages = req.Messages
	resp := s.responses[s.calls]
	s.calls++
	return resp, nil
}

func (s *scriptedLLM) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		resp, _ := s.Generate(ctx, req)
		if resp.Reasoning != "" {
			if !yield(llm.StreamChunk{ReasoningDelta: resp.Reasoning}, nil) {
				return
			}
		}
		// A response scripted with StopReason FinishThinkingBudget simulates
		// a provider that cut the round mid-reasoning: the budget chunk is
		// the last thing the stream produces.
		if resp.StopReason == llm.FinishThinkingBudget {
			yield(llm.StreamChunk{FinishReason: llm.FinishThinkingBudget}, nil)
			return
		}
		if resp.Content != "" {
			if !yield(llm.StreamChunk{Delta: resp.Content}, nil) {
				return
			}
		}
		if len(resp.ToolCalls) > 0 {
			if !yield(llm.StreamChunk{ToolCalls: resp.ToolCalls, FinishReason: "tool_calls"}, nil) {
				return
			}
		}
	}
}

func (s *scriptedLLM) Name() string                           { return "scripted" }
func (s *scriptedLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} }

func editTestRegistry(t *testing.T, executed *int) tools.Registry {
	t.Helper()
	reg := tools.NewRegistry()
	err := reg.Register(tools.Tool{
		Name:        "edit",
		Description: "edit",
		InputSchema: json.RawMessage(`{"type":"object"}`),
		Handler: func(_ context.Context, _ json.RawMessage) (tools.ToolResult, error) {
			*executed++
			return tools.ToolResult{Content: "ok"}, nil
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	return reg
}

// TestRunStream_RecoversFromUnparsedToolCall reproduces the failure seen
// live with Qwen3.6 + llama.cpp: the model writes its tool call as plain
// text inside its reasoning ("<tool_call><function=edit>...") so the server
// never parses it, the round has no tool calls, and the old loop simply
// ended the turn — abandoning the task mid-way with "now I'll fix X:" as the
// last words. The loop must instead nudge the model and let it re-issue the
// call for real.
func TestRunStream_RecoversFromUnparsedToolCall(t *testing.T) {
	executed := 0
	stub := &scriptedLLM{responses: []llm.CompletionResponse{
		// Round 1: tool call emitted as text inside reasoning — unparsed.
		{Reasoning: "I'll fix it now <tool_call> <function=edit> <parameter=path>x.py</parameter> </tool_call>", Content: "Voy a corregirlo:"},
		// Round 2 (after the nudge): a real, parsed tool call.
		{ToolCalls: []llm.ToolCall{{ID: "1", Name: "edit", Arguments: json.RawMessage(`{}`)}}},
		// Round 3: final answer.
		{Content: "Listo, corregido."},
	}}

	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})

	var final strings.Builder
	for chunk, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "arregla x.py"}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		final.WriteString(chunk.Delta)
	}

	if executed != 1 {
		t.Fatalf("expected the re-issued tool call to execute once, got %d", executed)
	}
	if !strings.Contains(final.String(), "Listo, corregido.") {
		t.Fatalf("expected the turn to continue to a final answer, got %q", final.String())
	}
	// The corrective nudge must have been sent to the model.
	foundNudge := false
	for _, m := range stub.lastMessages {
		if m.Role == llm.RoleUser && strings.Contains(m.Content, "NOT executed") {
			foundNudge = true
		}
	}
	if !foundNudge {
		t.Fatal("expected the corrective nudge in the follow-up request messages")
	}
}

// TestRunStream_NudgeGivesUpAfterLimit keeps a model that never emits a real
// tool call from ping-ponging forever: after maxUnparsedToolCallNudges the
// turn ends normally with whatever content there is.
func TestRunStream_NudgeGivesUpAfterLimit(t *testing.T) {
	executed := 0
	bad := llm.CompletionResponse{Content: "texto con <tool_call> falso"}
	stub := &scriptedLLM{responses: []llm.CompletionResponse{bad, bad, bad, bad}}

	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})

	rounds := 0
	for _, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "haz algo"}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
	}
	_ = rounds

	if stub.calls != maxUnparsedToolCallNudges+1 {
		t.Fatalf("expected %d rounds (original + nudges), got %d", maxUnparsedToolCallNudges+1, stub.calls)
	}
	if executed != 0 {
		t.Fatalf("no tool should have executed, got %d", executed)
	}
}

// TestRun_RecoversFromUnparsedToolCall covers the non-streaming path.
func TestRun_RecoversFromUnparsedToolCall(t *testing.T) {
	executed := 0
	stub := &scriptedLLM{responses: []llm.CompletionResponse{
		{Content: "ahora lo edito: <function=edit><parameter=path>x.py</parameter>"},
		{ToolCalls: []llm.ToolCall{{ID: "1", Name: "edit", Arguments: json.RawMessage(`{}`)}}},
		{Content: "Hecho."},
	}}

	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "arregla x.py"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if executed != 1 {
		t.Fatalf("expected the re-issued tool call to execute once, got %d", executed)
	}
	if resp.Content != "Hecho." {
		t.Fatalf("expected the final answer, got %q", resp.Content)
	}
}