package model

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// anthropicSeqServer answers successive /v1/messages calls with a fixed script
// of bodies (the last body repeats) and records every decoded request body, so a
// test can assert both the retry count and the retried max_tokens.
type anthropicSeqServer struct {
	*httptest.Server
	mu     sync.Mutex
	bodies []map[string]any
}

func newAnthropicSeqServer(t *testing.T, contentType string, bodies ...string) *anthropicSeqServer {
	t.Helper()
	s := &anthropicSeqServer{}
	idx := 0
	s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		_ = json.NewDecoder(r.Body).Decode(&body)
		s.mu.Lock()
		s.bodies = append(s.bodies, body)
		n := idx
		idx++
		s.mu.Unlock()
		if n >= len(bodies) {
			n = len(bodies) - 1
		}
		w.Header().Set("Content-Type", contentType)
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, bodies[n])
	}))
	t.Cleanup(s.Server.Close)
	return s
}

func (s *anthropicSeqServer) calls() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.bodies)
}

func (s *anthropicSeqServer) body(i int) map[string]any {
	s.mu.Lock()
	defer s.mu.Unlock()
	if i >= len(s.bodies) {
		return nil
	}
	return s.bodies[i]
}

// anthropicThinkingOnlyResponse is a thinking-only, budget-truncated reply: the
// reasoning consumed the whole max_tokens budget, so no text block was emitted.
const anthropicThinkingOnlyResponse = `{
  "id":"msg_think","model":"claude-test",
  "content":[{"type":"thinking","thinking":"SECRET-REASONING long long long"}],
  "stop_reason":"max_tokens",
  "usage":{"input_tokens":5,"output_tokens":1024}
}`

const anthropicEndTurnTextResponse = `{
  "id":"msg_text","model":"claude-test",
  "content":[{"type":"text","text":"答案"}],
  "stop_reason":"end_turn",
  "usage":{"input_tokens":5,"output_tokens":10}
}`

const anthropicToolUseResponse = `{
  "id":"msg_tool","model":"claude-test",
  "content":[{"type":"tool_use","id":"toolu_1","name":"datetime","input":{"tz":"Asia/Shanghai"}}],
  "stop_reason":"tool_use",
  "usage":{"input_tokens":5,"output_tokens":8}
}`

// anthropicTruncatedThinkingSSE is the streaming form of a thinking-only
// truncation: a thinking block, then message_delta stop_reason max_tokens.
const anthropicTruncatedThinkingSSE = `data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}

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

data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"SECRET-REASONING"}}

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

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

data: {"type":"message_stop"}

`

const anthropicEndTurnTextSSE = `data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}

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

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

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

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

data: {"type":"message_stop"}

`

// TestAnthropicChat_TruncatedThinking_RetriesWithDoubledBudget proves the
// primary recovery: a max_tokens truncation with only a thinking block retries
// exactly once with the token budget doubled and surfaces the retry's text.
func TestAnthropicChat_TruncatedThinking_RetriesWithDoubledBudget(t *testing.T) {
	t.Parallel()
	srv := newAnthropicSeqServer(t, "application/json", anthropicThinkingOnlyResponse, anthropicEndTurnTextResponse)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	out, err := ad.Chat(context.Background(), Request{
		Messages:  []Message{TextMessage(RoleUser, "hi")},
		MaxTokens: 1024,
	}, Options{})
	if err != nil {
		t.Fatalf("Chat error: %v", err)
	}
	if srv.calls() != 2 {
		t.Fatalf("requests = %d, want 2 (one retry)", srv.calls())
	}
	if got := srv.body(0)["max_tokens"]; got != float64(1024) {
		t.Errorf("first request max_tokens = %v, want 1024", got)
	}
	if got := srv.body(1)["max_tokens"]; got != float64(2048) {
		t.Errorf("retry request max_tokens = %v, want 2048 (doubled)", got)
	}
	if out.Content != "答案" {
		t.Errorf("Content = %q, want the retry's text", out.Content)
	}
	if out.StopReason != "end_turn" {
		t.Errorf("StopReason = %q, want end_turn", out.StopReason)
	}
	if strings.Contains(out.Content, "SECRET-REASONING") {
		t.Errorf("Content leaked thinking: %q", out.Content)
	}
}

// TestAnthropicChat_TruncatedThinking_RetryStillEmpty_FallsBackToHint proves the
// safety net: when the retry is truncated too, the caller gets a non-empty
// placeholder instead of a silent empty reply, and the thinking never leaks.
func TestAnthropicChat_TruncatedThinking_RetryStillEmpty_FallsBackToHint(t *testing.T) {
	t.Parallel()
	srv := newAnthropicSeqServer(t, "application/json", anthropicThinkingOnlyResponse, anthropicThinkingOnlyResponse)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	out, err := ad.Chat(context.Background(), Request{
		Messages:  []Message{TextMessage(RoleUser, "hi")},
		MaxTokens: 8192,
	}, Options{})
	if err != nil {
		t.Fatalf("Chat error: %v", err)
	}
	if srv.calls() != 2 {
		t.Fatalf("requests = %d, want 2 (retry once, never more)", srv.calls())
	}
	if got := srv.body(1)["max_tokens"]; got != float64(16384) {
		t.Errorf("retry request max_tokens = %v, want 16384", got)
	}
	if out.Content != truncationEmptyHint {
		t.Errorf("Content = %q, want the non-empty fallback hint", out.Content)
	}
	if out.Content == "" {
		t.Error("Content is empty: a truncated turn must never surface as a silent reply")
	}
	if strings.Contains(out.Content, "SECRET-REASONING") {
		t.Errorf("Content leaked thinking: %q", out.Content)
	}
	if out.Message.Role != RoleAssistant || len(out.Message.Content) != 1 || out.Message.Content[0].Text != truncationEmptyHint {
		t.Errorf("Message = %+v, want a single assistant hint part", out.Message)
	}
}

// TestAnthropicChat_NormalStop_NoRetry proves a normal end_turn answer and a
// native tool_use round (whose body is legitimately empty) are left untouched:
// no retry, no fallback hint.
func TestAnthropicChat_NormalStop_NoRetry(t *testing.T) {
	t.Parallel()
	t.Run("end_turn", func(t *testing.T) {
		t.Parallel()
		srv := newAnthropicSeqServer(t, "application/json", anthropicEndTurnTextResponse)
		ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
		out, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}, MaxTokens: 1024}, Options{})
		if err != nil {
			t.Fatalf("Chat error: %v", err)
		}
		if srv.calls() != 1 {
			t.Errorf("requests = %d, want 1 (no retry on end_turn)", srv.calls())
		}
		if out.Content != "答案" {
			t.Errorf("Content = %q, want 答案", out.Content)
		}
	})
	t.Run("tool_use", func(t *testing.T) {
		t.Parallel()
		srv := newAnthropicSeqServer(t, "application/json", anthropicToolUseResponse)
		ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
		out, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}, MaxTokens: 1024}, Options{})
		if err != nil {
			t.Fatalf("Chat error: %v", err)
		}
		if srv.calls() != 1 {
			t.Errorf("requests = %d, want 1 (a tool_use round is not a truncation)", srv.calls())
		}
		if len(out.ToolCalls) != 1 {
			t.Fatalf("ToolCalls = %+v, want 1", out.ToolCalls)
		}
		if out.Content != "" {
			t.Errorf("Content = %q, want empty (tool-only round)", out.Content)
		}
	})
}

// TestAnthropicChatStream_TruncatedThinking_RetriesWithDoubledBudget proves the
// streaming retry: the retry's thinking is suppressed (no duplicate block) while
// its text streams through, and the doubled budget is sent.
func TestAnthropicChatStream_TruncatedThinking_RetriesWithDoubledBudget(t *testing.T) {
	t.Parallel()
	srv := newAnthropicSeqServer(t, "text/event-stream", anthropicTruncatedThinkingSSE, anthropicEndTurnTextSSE)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	var events []StreamEvent
	out, err := ad.ChatStream(context.Background(), Request{
		Messages:  []Message{TextMessage(RoleUser, "hi")},
		MaxTokens: 1024,
	}, func(ev StreamEvent) { events = append(events, ev) })
	if err != nil {
		t.Fatalf("ChatStream error: %v", err)
	}
	if srv.calls() != 2 {
		t.Fatalf("requests = %d, want 2 (one retry)", srv.calls())
	}
	if got := srv.body(1)["max_tokens"]; got != float64(2048) {
		t.Errorf("retry request max_tokens = %v, want 2048", got)
	}
	if out.Content != "好" {
		t.Errorf("Content = %q, want the retry's text", out.Content)
	}
	if out.StopReason != "end_turn" {
		t.Errorf("StopReason = %q, want end_turn", out.StopReason)
	}
	want := []StreamEvent{
		{Kind: StreamKindThinkingStart, Index: 0},
		{Kind: StreamKindThinkingDelta, Index: 0, Delta: "SECRET-REASONING"},
		{Kind: StreamKindThinkingEnd, Index: 0},
		{Kind: StreamKindTextStart, Index: 0},
		{Kind: StreamKindTextDelta, Index: 0, Delta: "好"},
		{Kind: StreamKindTextEnd, Index: 0},
	}
	if !reflect.DeepEqual(events, want) {
		t.Errorf("events = %+v, want %+v (retry thinking suppressed, retry text streamed)", events, want)
	}
}

// TestAnthropicChatStream_TruncatedThinking_RetryStillEmpty_EmitsHint proves the
// streaming safety net: when the retry is truncated too, a non-empty placeholder
// is emitted as a text segment (and returned as the body) so the frontend is not
// left with a silent truncated turn.
func TestAnthropicChatStream_TruncatedThinking_RetryStillEmpty_EmitsHint(t *testing.T) {
	t.Parallel()
	srv := newAnthropicSeqServer(t, "text/event-stream", anthropicTruncatedThinkingSSE, anthropicTruncatedThinkingSSE)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	var events []StreamEvent
	out, err := ad.ChatStream(context.Background(), Request{
		Messages:  []Message{TextMessage(RoleUser, "hi")},
		MaxTokens: 1024,
	}, func(ev StreamEvent) { events = append(events, ev) })
	if err != nil {
		t.Fatalf("ChatStream error: %v", err)
	}
	if srv.calls() != 2 {
		t.Fatalf("requests = %d, want 2 (retry once, never more)", srv.calls())
	}
	if out.Content != truncationEmptyHint {
		t.Errorf("Content = %q, want the non-empty fallback hint", out.Content)
	}
	want := []StreamEvent{
		{Kind: StreamKindThinkingStart, Index: 0},
		{Kind: StreamKindThinkingDelta, Index: 0, Delta: "SECRET-REASONING"},
		{Kind: StreamKindThinkingEnd, Index: 0},
		{Kind: StreamKindTextStart, Index: 0},
		{Kind: StreamKindTextDelta, Index: 0, Delta: truncationEmptyHint},
		{Kind: StreamKindTextEnd, Index: 0},
	}
	if !reflect.DeepEqual(events, want) {
		t.Errorf("events = %+v, want %+v (hint streamed, thinking not re-emitted)", events, want)
	}
}

// TestTruncationBudget covers the budget arithmetic that bounds the retry.
func TestTruncationBudget(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name           string
		requested, def int
		wantEffective  int
		wantRetry      int
	}{
		{"explicit doubles", 1024, anthropicDefaultMaxTokens, 1024, 2048},
		{"default doubles", 0, anthropicDefaultMaxTokens, anthropicDefaultMaxTokens, 16384},
		{"capped", 20000, anthropicDefaultMaxTokens, 20000, truncationRetryCap},
		{"nothing to double", 0, 0, 0, 0},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if got := effectiveMaxTokens(tc.requested, tc.def); got != tc.wantEffective {
				t.Errorf("effectiveMaxTokens(%d, %d) = %d, want %d", tc.requested, tc.def, got, tc.wantEffective)
			}
			if got := retryMaxTokens(tc.requested, tc.def); got != tc.wantRetry {
				t.Errorf("retryMaxTokens(%d, %d) = %d, want %d", tc.requested, tc.def, got, tc.wantRetry)
			}
		})
	}
}

// TestOpenAIChat_StopReasonMapping asserts the OpenAI adapter fills the unified
// Response.StopReason from finish_reason on both the non-streaming and streaming
// paths, so a "length" truncation is identifiable across protocols.
func TestOpenAIChat_StopReasonMapping(t *testing.T) {
	t.Parallel()
	t.Run("non_streaming", func(t *testing.T) {
		t.Parallel()
		body := `{"id":"c1","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"cut"},"finish_reason":"length"}],"usage":{"prompt_tokens":5,"completion_tokens":64,"total_tokens":69}}`
		srv := newCompletionServer(t, http.StatusOK, "application/json", body, nil)
		ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"})
		out, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
		if err != nil {
			t.Fatalf("Chat error: %v", err)
		}
		if out.StopReason != "length" {
			t.Errorf("StopReason = %q, want length", out.StopReason)
		}
	})
	t.Run("streaming", func(t *testing.T) {
		t.Parallel()
		sse := strings.Join([]string{
			`data: {"choices":[{"index":0,"delta":{"content":"Hi"}}]}`,
			``,
			`data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`,
			``,
			`data: [DONE]`,
			``,
		}, "\n")
		srv := newCompletionServer(t, http.StatusOK, "text/event-stream", sse, nil)
		ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"})
		out, err := ad.ChatStream(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, nil)
		if err != nil {
			t.Fatalf("ChatStream error: %v", err)
		}
		if out.StopReason != "stop" {
			t.Errorf("StopReason = %q, want stop", out.StopReason)
		}
	})
}