package runtime

// This file implements the faux provider (mirrors pi providers/faux.ts) and the
// loop integration tests that drive the whole agent loop through it — the
// project's primary and only core test seam (US-002 / Testing Decisions, #16).
//
// Unlike loop_test.go, which drives the loop with a coarse StreamFn that emits
// only a terminal StreamDoneEvent, the faux provider is a real Provider whose
// StreamCompletion replays a *fine-grained* script of AssistantMessageEvents
// (start → text/toolcall deltas → done) — one scripted turn per call. It is
// wired into the loop via StreamFnFromProvider, the real seam, so the whole
// path (message_start / message_update / message_end deltas, the six hooks,
// truncation protection, parallel ordering, and EventStream cancellation) is
// covered end to end without mocking any loop-internal function.

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

	"github.com/smallnest/pigo/internal/agentcore"
	"github.com/smallnest/pigo/internal/agenttool"
	"github.com/smallnest/pigo/internal/provider"
)

// fauxTurn is one scripted assistant turn: the fine-grained stream events the
// faux provider replays for a single StreamCompletion call.
type fauxTurn []provider.AssistantMessageEvent

// textTurn scripts a turn that streams text as start → text delta → done(end_turn).
func textTurn(text string) fauxTurn {
	partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}
	withText := partial
	withText.Content = agentcore.ContentList{agentcore.NewTextContent(text)}
	final := withText
	final.StopReason = agentcore.StopReasonEndTurn
	return fauxTurn{
		provider.StreamStartEvent{Partial: partial},
		provider.StreamTextEvent{Partial: withText},
		provider.StreamDoneEvent{Message: final},
	}
}

// toolCallTurn scripts a turn that streams one tool call as
// start → toolcall delta → done(tool_use).
func toolCallTurn(id, name, args string) fauxTurn {
	partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}
	withCall := partial
	withCall.Content = agentcore.ContentList{agentcore.NewToolCallContent(id, name, json.RawMessage(args))}
	final := withCall
	final.StopReason = agentcore.StopReasonToolUse
	return fauxTurn{
		provider.StreamStartEvent{Partial: partial},
		provider.StreamToolCallEvent{Partial: withCall},
		provider.StreamDoneEvent{Message: final},
	}
}

// fauxProvider is a real Provider that replays one scripted turn per
// StreamCompletion call, in order. It records every request it received so
// tests can assert what the loop actually sent (model, context, config). Once
// the script is exhausted it replays a plain end_turn turn.
type fauxProvider struct {
	name   string
	models []provider.Model
	turns  []fauxTurn

	mu       sync.Mutex
	calls    int
	requests []provider.CompletionRequest
	// delay optionally slows each delta emit, used by the cancellation test to
	// keep the stream open long enough to cancel mid-flight.
	delay time.Duration
}

func (p *fauxProvider) Name() string             { return p.name }
func (p *fauxProvider) Models() []provider.Model { return p.models }

func (p *fauxProvider) callCount() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.calls
}

func (p *fauxProvider) requestAt(i int) provider.CompletionRequest {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.requests[i]
}

func (p *fauxProvider) StreamCompletion(ctx context.Context, req provider.CompletionRequest) (*provider.AssistantMessageEventStream, error) {
	p.mu.Lock()
	idx := p.calls
	p.calls++
	p.requests = append(p.requests, req)
	var turn fauxTurn
	if idx < len(p.turns) {
		turn = p.turns[idx]
	} else {
		turn = textTurn("")
	}
	delay := p.delay
	p.mu.Unlock()

	s := provider.NewAssistantMessageEventStream(0)
	go func() {
		for _, ev := range turn {
			if delay > 0 {
				select {
				case <-time.After(delay):
				case <-ctx.Done():
					s.SetError(ctx.Err())
					s.Close()
					return
				}
			}
			if err := s.Emit(ctx, ev); err != nil {
				s.SetError(err)
				s.Close()
				return
			}
		}
		s.Close()
	}()
	return s, nil
}

// newFauxRunCfg wires a faux provider into the loop via StreamFnFromProvider
// (the real seam) and registers the given tools. No loop-internal function is
// mocked — only the provider boundary.
func newFauxRunCfg(p *fauxProvider, tools ...agentcore.AgentTool) RunConfig {
	reg := agenttool.NewToolRegistry()
	for _, tl := range tools {
		_ = reg.Register(tl)
	}
	return RunConfig{
		LoopConfig: LoopConfig{Model: "faux", Stream: provider.StreamFnFromProvider(p)},
		Batch:      agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}},
	}
}

// TestFauxProviderTextToolText drives the flagship seam scenario end to end:
// text → tool call → text over the real loop, asserting both the AgentEvent
// stream shape and the final []AgentMessage. Nothing loop-internal is mocked.
func TestFauxProviderTextToolText(t *testing.T) {
	p := &fauxProvider{
		name:   "faux",
		models: []provider.Model{{Provider: "faux", ID: "faux"}},
		turns: []fauxTurn{
			textTurn("thinking about it"),                     // turn 1: plain text, no tool
			toolCallTurn("call-1", "echo", `{"msg":"hello"}`), // turn 2: tool call
			textTurn("all done"),                              // turn 3: final text
		},
	}
	// GetFollowUpMessages injects a follow-up once so the loop advances past the
	// first natural (text-only) turn end into the tool-call turn.
	served := false
	cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false))
	cfg.GetFollowUpMessages = func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage {
		if served {
			return nil
		}
		served = true
		return []agentcore.AgentMessage{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("go on")}}}
	}
	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}}

	kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg))

	// Event shape: message deltas must appear (start/update/end), a tool
	// executed exactly once, and the run bookended by agent_start/agent_end.
	if kinds[0] != agentcore.EventAgentStart || kinds[len(kinds)-1] != agentcore.EventAgentEnd {
		t.Fatalf("run must be bracketed by agent_start/agent_end, got %v", kinds)
	}
	if countKind(kinds, agentcore.EventMessageStart) < 3 || countKind(kinds, agentcore.EventMessageEnd) < 3 {
		t.Errorf("expected fine-grained message deltas for each turn, got %v", kinds)
	}
	if countKind(kinds, agentcore.EventMessageUpdate) < 3 {
		t.Errorf("expected message_update deltas (text/toolcall), got %v", kinds)
	}
	if got := countKind(kinds, agentcore.EventToolExecutionStart); got != 1 {
		t.Errorf("expected 1 tool_execution_start, got %d in %v", got, kinds)
	}
	if got := countKind(kinds, agentcore.EventToolExecutionEnd); got != 1 {
		t.Errorf("expected 1 tool_execution_end, got %d in %v", got, kinds)
	}
	if got := countKind(kinds, agentcore.EventTurnStart); got != 3 {
		t.Errorf("expected 3 turns (text→tool→text), got %d in %v", got, kinds)
	}

	// Final messages: assistant(text) + user(follow-up) + assistant(tool) +
	// toolResult + assistant(text) = 5, in order.
	if len(msgs) != 5 {
		t.Fatalf("expected 5 new messages, got %d: %+v", len(msgs), msgs)
	}
	if a, ok := msgs[0].(agentcore.AssistantMessage); !ok || textContentOf(a.Content) != "thinking about it" {
		t.Errorf("msg[0] should be the first text assistant message, got %T %+v", msgs[0], msgs[0])
	}
	if _, ok := msgs[1].(agentcore.UserMessage); !ok {
		t.Errorf("msg[1] should be the injected follow-up user message, got %T", msgs[1])
	}
	if a, ok := msgs[2].(agentcore.AssistantMessage); !ok || len(a.ToolCalls()) != 1 {
		t.Errorf("msg[2] should be the tool-call assistant message, got %T %+v", msgs[2], msgs[2])
	}
	tr, ok := msgs[3].(agentcore.ToolResultMessage)
	if !ok || tr.ToolCallID != "call-1" || tr.IsError {
		t.Errorf("msg[3] should be the successful echo tool result, got %T %+v", msgs[3], msgs[3])
	}
	if a, ok := msgs[4].(agentcore.AssistantMessage); !ok || textContentOf(a.Content) != "all done" {
		t.Errorf("msg[4] should be the final text assistant message, got %T %+v", msgs[4], msgs[4])
	}

	// The loop must have driven the provider exactly three times, each carrying
	// the growing context and the configured model.
	if p.callCount() != 3 {
		t.Fatalf("provider called %d times, want 3", p.callCount())
	}
	if req := p.requestAt(0); req.Model != "faux" {
		t.Errorf("provider request model = %q, want faux", req.Model)
	}
}

// textContentOf returns the concatenated text of a content list.
func textContentOf(list agentcore.ContentList) string {
	var s string
	for _, c := range list {
		if tc, ok := c.(agentcore.TextContent); ok {
			s += tc.Text
		}
	}
	return s
}

// TestFauxSeamSixHooks exercises all six loop hooks through the real seam in a
// single run: the two per-request LoopConfig hooks (TransformContext,
// ConvertToLlm resolved via GetAPIKey) and the four RunConfig hooks
// (GetFollowUpMessages, GetSteeringMessages, PrepareNextTurn,
// ShouldStopAfterTurn). Each hook records that it fired and, where observable,
// that its effect reached the provider request.
func TestFauxSeamSixHooks(t *testing.T) {
	p := &fauxProvider{
		name:   "faux",
		models: []provider.Model{{Provider: "faux", ID: "faux"}},
		turns: []fauxTurn{
			toolCallTurn("call-1", "echo", `{}`), // turn 1: tool → afterTurn hooks fire
			textTurn("second"),                   // turn 2: end (after model swap)
		},
	}
	var fired struct {
		transform, convert, apiKey, followUp, steering, prepare, shouldStop bool
	}
	swapped := "swapped-model"
	cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false))
	cfg.Provider = "faux"
	cfg.TransformContext = func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList {
		fired.transform = true
		return msgs
	}
	cfg.ConvertToLlm = func(msgs agentcore.MessageList) agentcore.MessageList {
		fired.convert = true
		return msgs
	}
	cfg.GetAPIKey = func(ctx context.Context, provider string) string {
		fired.apiKey = true
		return "dyn-key"
	}
	cfg.GetSteeringMessages = func(ctx context.Context) []agentcore.AgentMessage {
		fired.steering = true
		return nil
	}
	cfg.PrepareNextTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) *TurnUpdate {
		fired.prepare = true
		return &TurnUpdate{Model: &swapped}
	}
	stopCalls := 0
	cfg.ShouldStopAfterTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) bool {
		fired.shouldStop = true
		stopCalls++
		return false // never stop early; let the run end naturally
	}
	cfg.GetFollowUpMessages = func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage {
		fired.followUp = true
		// No follow-up: the tool-call turn already drives turn 2, so the run
		// ends naturally after the second turn.
		return nil
	}
	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}}

	collectStream(t, agentLoop(context.Background(), agentCtx, cfg))

	if !fired.transform || !fired.convert || !fired.apiKey {
		t.Errorf("per-request hooks not all fired: %+v", fired)
	}
	if !fired.followUp || !fired.steering || !fired.prepare || !fired.shouldStop {
		t.Errorf("per-turn hooks not all fired: %+v", fired)
	}
	if stopCalls == 0 {
		t.Error("ShouldStopAfterTurn was never consulted")
	}
	// GetAPIKey's dynamic key must have reached the provider request config.
	if got := p.requestAt(0).Config.APIKey; got != "dyn-key" {
		t.Errorf("GetAPIKey result not threaded to provider, APIKey = %q", got)
	}
	// PrepareNextTurn swapped the model before turn 2.
	if p.callCount() >= 2 {
		if got := p.requestAt(1).Config.APIKey; got != "dyn-key" {
			t.Errorf("turn 2 APIKey = %q, want dyn-key", got)
		}
		if got := p.requestAt(1).Model; got != swapped {
			t.Errorf("PrepareNextTurn model swap not applied, turn 2 model = %q, want %q", got, swapped)
		}
	}
}

// TestFauxSeamTruncationProtection verifies that a truncated (stopReason=length)
// tool-call turn is protected: the tool is NOT executed and a synthesized failed
// tool result is fed back, all through the seam.
func TestFauxSeamTruncationProtection(t *testing.T) {
	// Turn 1: a tool call that arrives truncated. Turn 2: end.
	truncPartial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("t1", "echo", json.RawMessage(`{}`))}}
	truncFinal := truncPartial
	truncFinal.StopReason = agentcore.StopReasonLength
	p := &fauxProvider{
		name: "faux",
		turns: []fauxTurn{
			{
				provider.StreamStartEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}},
				provider.StreamToolCallEvent{Partial: truncPartial},
				provider.StreamDoneEvent{Message: truncFinal},
			},
			textTurn("recovered"),
		},
	}
	cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false))
	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}}

	kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg))

	if countKind(kinds, agentcore.EventToolExecutionEnd) != 0 {
		t.Errorf("truncated tool call must not execute, got %v", kinds)
	}
	var foundFail bool
	for _, m := range msgs {
		if tr, ok := m.(agentcore.ToolResultMessage); ok && tr.IsError && tr.ToolCallID == "t1" {
			foundFail = true
		}
	}
	if !foundFail {
		t.Errorf("expected a synthesized failed tool result for the truncated call, got %+v", msgs)
	}
}

// TestFauxSeamParallelOrderingPreserved verifies that a turn with multiple
// parallel tool calls yields tool results in source order regardless of which
// tool finishes first, driven through the seam.
func TestFauxSeamParallelOrderingPreserved(t *testing.T) {
	// One assistant turn with three tool calls in a fixed order; the tools sleep
	// in reverse so completion order differs from source order.
	partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{
		agentcore.NewToolCallContent("a0", "slow", json.RawMessage(`{}`)),
		agentcore.NewToolCallContent("a1", "mid", json.RawMessage(`{}`)),
		agentcore.NewToolCallContent("a2", "fast", json.RawMessage(`{}`)),
	}}
	final := partial
	final.StopReason = agentcore.StopReasonToolUse
	p := &fauxProvider{
		turns: []fauxTurn{
			{provider.StreamStartEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}, provider.StreamToolCallEvent{Partial: partial}, provider.StreamDoneEvent{Message: final}},
			textTurn("done"),
		},
	}
	mk := func(name string, delay time.Duration) execTool {
		return execTool{
			name: name,
			mode: agentcore.ToolExecutionParallel,
			run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) {
				time.Sleep(delay)
				return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil
			},
		}
	}
	cfg := newFauxRunCfg(p, mk("slow", 25*time.Millisecond), mk("mid", 12*time.Millisecond), mk("fast", 1*time.Millisecond))
	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}}

	_, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg))

	var order []string
	for _, m := range msgs {
		if tr, ok := m.(agentcore.ToolResultMessage); ok {
			order = append(order, tr.ToolCallID)
		}
	}
	want := []string{"a0", "a1", "a2"}
	if len(order) != 3 || order[0] != want[0] || order[1] != want[1] || order[2] != want[2] {
		t.Errorf("parallel tool results out of source order: got %v, want %v", order, want)
	}
}

// TestFauxSeamStreamCancellation verifies that cancelling the context stops the
// run: the consumer stops receiving events and Result reports the cancellation,
// exercised through the seam with a provider that streams slowly.
func TestFauxSeamStreamCancellation(t *testing.T) {
	p := &fauxProvider{
		turns: []fauxTurn{textTurn("never fully delivered")},
		delay: 50 * time.Millisecond, // slow enough to cancel mid-stream
	}
	cfg := newFauxRunCfg(p)
	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}}

	ctx, cancel := context.WithCancel(context.Background())
	s := agentLoop(ctx, agentCtx, cfg)

	// Read the first event, then cancel while the provider is still streaming.
	<-s.Events()
	cancel()
	// Drain remaining events (must terminate, not hang).
	for range s.Events() {
	}
	if _, err := s.Result(context.Background()); err != nil {
		// A set result is also acceptable (the run may have finished emitting
		// agent_end before cancellation propagated); but if an error is set it
		// must be the cancellation.
		if err != context.Canceled {
			t.Errorf("cancelled run result error = %v, want context.Canceled or nil", err)
		}
	}
}