package provider

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

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

// A recorded Anthropic Messages API SSE stream covering a text block, a
// thinking block (with signature), and a tool_use block, ending with a
// tool_use stop reason and output-token usage. Trimmed but structurally
// faithful to the real wire format.
const anthropicToolUseSSE = `event: message_start
data: {"type":"message_start","message":{"id":"msg_01ABC","model":"claude-opus-4-8","usage":{"input_tokens":42,"output_tokens":1}}}

event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Let me think. "}}

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Use the tool."}}

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sigABC"}}

event: content_block_stop
data: {"type":"content_block_stop","index":0}

event: content_block_start
data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}

event: content_block_delta
data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"I'll check "}}

event: content_block_delta
data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"the weather."}}

event: content_block_stop
data: {"type":"content_block_stop","index":1}

event: content_block_start
data: {"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"toolu_01","name":"get_weather"}}

event: content_block_delta
data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}}

event: content_block_delta
data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":" \"SF\"}"}}

event: content_block_stop
data: {"type":"content_block_stop","index":2}

event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":57}}

event: message_stop
data: {"type":"message_stop"}

`

// feedSSE splits a recorded SSE body into events (blank-line separated),
// extracts each `data:` payload, and drives the decoder exactly as the
// transport pump would, returning all emitted events plus the final message.
func feedSSE(t *testing.T, dec Decoder, body string) ([]StreamEvent, agentcore.AssistantMessage) {
	t.Helper()
	var events []StreamEvent
	for _, block := range strings.Split(body, "\n\n") {
		var payload strings.Builder
		for _, line := range strings.Split(block, "\n") {
			line = strings.TrimRight(line, "\r")
			if strings.HasPrefix(line, "data:") {
				payload.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
			}
		}
		if payload.Len() == 0 {
			continue
		}
		// [DONE] is a transport-level terminator, not a decoder payload.
		if payload.String() == "[DONE]" {
			continue
		}
		evs, err := dec.Decode([]byte(payload.String()))
		if err != nil {
			t.Fatalf("decode %q: %v", payload.String(), err)
		}
		events = append(events, evs...)
	}
	finalEvents, err := dec.Finish()
	if err != nil {
		t.Fatalf("finish: %v", err)
	}
	events = append(events, finalEvents...)

	var final agentcore.AssistantMessage
	for _, ev := range events {
		if d, ok := ev.(StreamDoneEvent); ok {
			final = d.Message
		}
	}
	return events, final
}

func TestAnthropicDecoderToolUseStream(t *testing.T) {
	dec := NewAnthropicDecoder()
	events, final := feedSSE(t, dec, anthropicToolUseSSE)

	// The first emitted event must be a start event.
	if len(events) == 0 || events[0].EventKind() != StreamEventStart {
		t.Fatalf("expected a start event first, got %v", eventKinds(events))
	}
	// The last emitted event must be the terminal done event.
	if events[len(events)-1].EventKind() != StreamEventDone {
		t.Fatalf("expected a done event last, got %v", eventKinds(events))
	}

	// Stop reason: tool_use.
	if final.StopReason != agentcore.StopReasonToolUse {
		t.Errorf("stop reason = %q, want tool_use", final.StopReason)
	}
	// Usage: input from message_start, output from message_delta.
	if final.Usage == nil || final.Usage.InputTokens != 42 || final.Usage.OutputTokens != 57 {
		t.Errorf("usage = %+v, want input=42 output=57", final.Usage)
	}
	// Response identity from message_start.
	if final.ResponseID != "msg_01ABC" || final.ResponseModel != "claude-opus-4-8" {
		t.Errorf("response id/model = %q/%q", final.ResponseID, final.ResponseModel)
	}

	// Content blocks in index order: thinking, text, tool_use.
	if len(final.Content) != 3 {
		t.Fatalf("expected 3 content blocks, got %d: %+v", len(final.Content), final.Content)
	}
	th, ok := final.Content[0].(agentcore.ThinkingContent)
	if !ok || th.Thinking != "Let me think. Use the tool." {
		t.Errorf("thinking block = %+v", final.Content[0])
	}
	if th.ThinkingSignature != "sigABC" {
		t.Errorf("thinking signature = %q, want sigABC", th.ThinkingSignature)
	}
	txt, ok := final.Content[1].(agentcore.TextContent)
	if !ok || txt.Text != "I'll check the weather." {
		t.Errorf("text block = %+v", final.Content[1])
	}
	tool, ok := final.Content[2].(agentcore.ToolCallContent)
	if !ok || tool.Name != "get_weather" || tool.ID != "toolu_01" {
		t.Fatalf("tool block = %+v", final.Content[2])
	}
	// tool_use JSON must have accumulated into valid arguments.
	var args map[string]string
	if err := json.Unmarshal(tool.Arguments, &args); err != nil {
		t.Fatalf("tool arguments not valid JSON %q: %v", tool.Arguments, err)
	}
	if args["city"] != "SF" {
		t.Errorf("tool arguments = %v, want city=SF", args)
	}
}

func TestAnthropicDecoderTextOnlyEndTurn(t *testing.T) {
	body := `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-x","usage":{"input_tokens":10,"output_tokens":0}}}

data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}

data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}

data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}}

data: {"type":"content_block_stop","index":0}

data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}

data: {"type":"message_stop"}

`
	dec := NewAnthropicDecoder()
	_, final := feedSSE(t, dec, body)
	if final.StopReason != agentcore.StopReasonEndTurn {
		t.Errorf("stop reason = %q, want end_turn", final.StopReason)
	}
	if len(final.Content) != 1 {
		t.Fatalf("expected 1 content block, got %d", len(final.Content))
	}
	txt, ok := final.Content[0].(agentcore.TextContent)
	if !ok || txt.Text != "Hello world" {
		t.Errorf("text = %+v", final.Content[0])
	}
}

func TestAnthropicDecoderMaxTokensMapsToLength(t *testing.T) {
	body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}}

data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}

data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"truncated"}}

data: {"type":"content_block_stop","index":0}

data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":100}}

data: {"type":"message_stop"}

`
	dec := NewAnthropicDecoder()
	_, final := feedSSE(t, dec, body)
	if final.StopReason != agentcore.StopReasonLength {
		t.Errorf("max_tokens must map to length, got %q", final.StopReason)
	}
}

// TestAnthropicDecoderErrorEvent verifies an Anthropic `error` event becomes a
// decode error (which the transport turns into a terminal error event), never
// a panic.
func TestAnthropicDecoderErrorEvent(t *testing.T) {
	dec := NewAnthropicDecoder()
	_, err := dec.Decode([]byte(`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`))
	if err == nil {
		t.Fatal("error event must return a decode error")
	}
	if !strings.Contains(err.Error(), "overloaded_error") {
		t.Errorf("error should name the type, got %v", err)
	}
}

// TestAnthropicDecoderMalformedPayload verifies invalid JSON is a returned
// error (rides the stream as terminal error), not a panic.
func TestAnthropicDecoderMalformedPayload(t *testing.T) {
	dec := NewAnthropicDecoder()
	if _, err := dec.Decode([]byte(`{not json`)); err == nil {
		t.Fatal("malformed payload must return an error")
	}
}

// TestAnthropicDecoderFinishFlushesPartial verifies a stream cut short (no
// message_stop) still yields a done event on Finish so the partial isn't lost.
func TestAnthropicDecoderFinishFlushesPartial(t *testing.T) {
	body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}}

data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}

data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}

`
	dec := NewAnthropicDecoder()
	events, final := feedSSE(t, dec, body)
	if events[len(events)-1].EventKind() != StreamEventDone {
		t.Fatalf("Finish must emit a terminal done event, got %v", eventKinds(events))
	}
	// No message_delta arrived → default end_turn.
	if final.StopReason != agentcore.StopReasonEndTurn {
		t.Errorf("cut-short stream should default to end_turn, got %q", final.StopReason)
	}
	if len(final.Content) != 1 {
		t.Fatalf("expected the partial text block, got %+v", final.Content)
	}
}

// TestAnthropicDecoderThroughTransport wires the decoder through the real
// transport pump against a recorded SSE server, exercising the full path.
func TestAnthropicDecoderThroughTransport(t *testing.T) {
	srv := sseServer(t, anthropicToolUseSSE)
	defer srv.Close()

	stream, err := StreamRequest(context.Background(), TransportConfig{
		NewRequest: newReqFn(srv.URL),
		Decoder:    NewAnthropicDecoder(),
	})
	if err != nil {
		t.Fatalf("StreamRequest: %v", err)
	}
	var kinds []string
	for ev := range stream.Events() {
		kinds = append(kinds, ev.EventKind())
	}
	final, resErr := stream.Result(context.Background())
	if resErr != nil {
		t.Fatalf("result: %v", resErr)
	}
	if final.StopReason != agentcore.StopReasonToolUse {
		t.Errorf("stop reason via transport = %q, want tool_use", final.StopReason)
	}
	if len(final.Content) != 3 {
		t.Errorf("expected 3 content blocks via transport, got %d", len(final.Content))
	}
	if kinds[len(kinds)-1] != StreamEventDone {
		t.Errorf("stream must end with done, got %v", kinds)
	}
}

func eventKinds(events []StreamEvent) []string {
	out := make([]string, len(events))
	for i, ev := range events {
		out[i] = ev.EventKind()
	}
	return out
}