package proxy_test

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"git.misaka.ren/M1saka/token_thief/config"
	"git.misaka.ren/M1saka/token_thief/logger"
	"git.misaka.ren/M1saka/token_thief/proxy"
)

type captureSubmitter struct {
	mu      sync.Mutex
	entries []*logger.LogEntry
}

func (c *captureSubmitter) Submit(e *logger.LogEntry) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries = append(c.entries, e)
}

func (c *captureSubmitter) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.entries)
}

func (c *captureSubmitter) Entry(i int) *logger.LogEntry {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.entries[i]
}

// fakeOpenAIStreamUpstream 模拟一个 OpenAI 兼容的 SSE 上游：
// 分 5 次往响应里 write 一行 SSE 数据，每次都 Flush。
func fakeOpenAIStreamUpstream() *httptest.Server {
	chunks := []string{
		`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"你"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"好"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
		"data: [DONE]\n\n",
	}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(http.StatusOK)
		flusher := w.(http.Flusher)
		for _, ch := range chunks {
			_, _ = io.WriteString(w, ch)
			flusher.Flush()
			time.Sleep(5 * time.Millisecond)
		}
	}))
}

func fakeAnthropicStreamUpstream() *httptest.Server {
	chunks := []string{
		`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-1","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
		`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n",
		`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"你"}}` + "\n\n",
		`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"好"}}` + "\n\n",
		`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n",
		`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
	}
	return newSSEUpstream(chunks)
}

func fakeGeminiStreamUpstream() *httptest.Server {
	chunks := []string{
		`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"index":0}]}` + "\n\n",
		`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
	}
	return newSSEUpstream(chunks)
}

func fakeUnknownStreamUpstream() *httptest.Server {
	return newSSEUpstream([]string{
		`event: custom` + "\n" + `data: not-json` + "\n\n",
	})
}

func fakeOpenAIStreamUpstreamThatStaysOpen(release <-chan struct{}) *httptest.Server {
	chunks := []string{
		`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
		"data: [DONE]\n\n",
	}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		flusher := w.(http.Flusher)
		for _, ch := range chunks {
			_, _ = io.WriteString(w, ch)
			flusher.Flush()
		}
		<-release
	}))
}

func newSSEUpstream(chunks []string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(http.StatusOK)
		flusher := w.(http.Flusher)
		for _, ch := range chunks {
			_, _ = io.WriteString(w, ch)
			flusher.Flush()
			time.Sleep(5 * time.Millisecond)
		}
	}))
}

func TestSSEChunkAssembly(t *testing.T) {
	upstream := fakeOpenAIStreamUpstream()
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 1024*1024)

	proxySrv := httptest.NewServer(h)
	defer proxySrv.Close()

	// 客户端走原始 TCP，逐字节读 + 打印，验证流式实时到达
	pu, _ := url.Parse(proxySrv.URL)
	conn, err := net.DialTimeout("tcp", pu.Host, 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	fmt.Fprintf(conn, "POST /v1/chat/completions HTTP/1.1\r\nHost: %s\r\nContent-Length: 0\r\n\r\n", pu.Host)
	_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))

	br := bufio.NewReader(conn)
	resp, err := http.ReadResponse(br, nil)
	if err != nil {
		t.Fatal(err)
	}
	clientBody, _ := io.ReadAll(resp.Body)

	// 等待 proxy.ServeHTTP 返回并把 entry 提交
	deadline := time.Now().Add(2 * time.Second)
	for sub.Len() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if sub.Len() == 0 {
		t.Fatal("no log entry captured")
	}
	e := sub.Entry(0)

	t.Logf("\n========== 客户端收到的字节 ==========\n%s", clientBody)
	t.Logf("\n========== 数据库 response_body 字段（按字节原样存储）==========\n%s", e.ResponseBody)
	t.Logf("\n========== 元数据 ==========")
	t.Logf("is_stream      = %v", e.IsStream)
	t.Logf("status_code    = %d", e.StatusCode)
	t.Logf("len(body)      = %d bytes", len(e.ResponseBody))
	t.Logf("response_truncated = %v", e.ResponseTruncated)

	if !strings.Contains(string(clientBody), `"content":"你"`) ||
		!strings.Contains(string(clientBody), `"content":"好"`) ||
		!strings.Contains(string(clientBody), "[DONE]") {
		t.Errorf("client response body 缺少预期 chunk 内容")
	}

	var captured struct {
		Choices []struct {
			Message struct {
				Role    string `json:"role"`
				Content string `json:"content"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("stream response_body should be assembled JSON: %v; body=%q", err, e.ResponseBody)
	}
	if len(captured.Choices) != 1 {
		t.Fatalf("assembled JSON choices length=%d, want 1", len(captured.Choices))
	}
	if captured.Choices[0].Message.Role != "assistant" {
		t.Errorf("assembled role=%q, want assistant", captured.Choices[0].Message.Role)
	}
	if captured.Choices[0].Message.Content != "你好" {
		t.Errorf("assembled content=%q, want 你好", captured.Choices[0].Message.Content)
	}
	if captured.Choices[0].FinishReason != "stop" {
		t.Errorf("assembled finish_reason=%q, want stop", captured.Choices[0].FinishReason)
	}
	if !e.IsStream {
		t.Errorf("is_stream 应为 true")
	}
}

func TestOpenAIStreamPreservesReasoningAndMetadata(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"think "},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"reasoning_content":"hard"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"content":"final"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{},"finish_reason":"stop","native_finish_reason":"stop"}]}` + "\n\n",
		"data: [DONE]\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/chat/completions")

	var captured struct {
		ID      string `json:"id"`
		Object  string `json:"object"`
		Created int64  `json:"created"`
		Model   string `json:"model"`
		Choices []struct {
			Index   int `json:"index"`
			Message struct {
				Role             string `json:"role"`
				Content          string `json:"content"`
				ReasoningContent string `json:"reasoning_content"`
			} `json:"message"`
			FinishReason       string `json:"finish_reason"`
			NativeFinishReason string `json:"native_finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("openai stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
	}
	if captured.ID != "resp-1" || captured.Object != "chat.completion" || captured.Created != 1779335544 || captured.Model != "gpt-5.4-mini-2026-03-17" {
		t.Fatalf("unexpected metadata: %+v", captured)
	}
	if len(captured.Choices) != 1 {
		t.Fatalf("choices length=%d, want 1", len(captured.Choices))
	}
	choice := captured.Choices[0]
	if choice.Message.Role != "assistant" || choice.Message.Content != "final" || choice.Message.ReasoningContent != "think hard" {
		t.Fatalf("unexpected message: %+v", choice.Message)
	}
	if choice.FinishReason != "stop" || choice.NativeFinishReason != "stop" {
		t.Fatalf("unexpected finish reasons: %+v", choice)
	}
}

func TestOpenAIStreamDoneDoesNotSubmitBeforeUpstreamCloses(t *testing.T) {
	release := make(chan struct{})
	upstream := fakeOpenAIStreamUpstreamThatStaysOpen(release)
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 1024*1024)

	proxySrv := httptest.NewServer(h)
	released := false
	defer func() {
		if !released {
			close(release)
		}
		proxySrv.Close()
		upstream.Close()
	}()

	clientDone := make(chan error, 1)
	go func() {
		resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
		if err != nil {
			clientDone <- err
			return
		}
		_, _ = io.ReadAll(resp.Body)
		_ = resp.Body.Close()
		clientDone <- nil
	}()

	time.Sleep(100 * time.Millisecond)
	if sub.Len() != 0 {
		t.Fatal("terminal SSE event must not submit before ReverseProxy returns")
	}
	close(release)
	released = true

	select {
	case err := <-clientDone:
		if err != nil {
			t.Fatal(err)
		}
		if sub.Len() != 1 {
			t.Fatalf("stream should be submitted once, got %d entries", sub.Len())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("client did not finish after upstream closed")
	}
}

func TestOpenAIStreamPreservesUsageChunk(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}` + "\n\n",
		`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}` + "\n\n",
		"data: [DONE]\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/chat/completions")

	var captured struct {
		ID                string `json:"id"`
		SystemFingerprint string `json:"system_fingerprint"`
		Choices           []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
			TotalTokens      int `json:"total_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("openai stream with usage should be assembled JSON: %v; body=%q", err, e.ResponseBody)
	}
	if captured.ID != "resp-usage" || captured.SystemFingerprint != "fp_123" {
		t.Fatalf("metadata not preserved: %+v", captured)
	}
	if len(captured.Choices) != 1 || captured.Choices[0].Message.Content != "ok" {
		t.Fatalf("choices not assembled: %+v", captured.Choices)
	}
	if captured.Usage.PromptTokens != 5 || captured.Usage.CompletionTokens != 2 || captured.Usage.TotalTokens != 7 {
		t.Fatalf("usage not preserved: %+v", captured.Usage)
	}
}

func TestOpenAIStreamAssemblesToolCalls(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{\"q\":"}}]},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"weather\"}"}}]},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}` + "\n\n",
		"data: [DONE]\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/chat/completions")

	var captured struct {
		Choices []struct {
			Message struct {
				Role      string `json:"role"`
				Content   string `json:"content"`
				ToolCalls []struct {
					ID       string `json:"id"`
					Type     string `json:"type"`
					Function struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("openai tool stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
	}
	if len(captured.Choices) != 1 || captured.Choices[0].FinishReason != "tool_calls" {
		t.Fatalf("unexpected choices: %+v", captured.Choices)
	}
	message := captured.Choices[0].Message
	if message.Role != "assistant" || message.Content != "" {
		t.Fatalf("unexpected message basics: %+v", message)
	}
	if len(message.ToolCalls) != 1 {
		t.Fatalf("tool_calls length=%d, want 1; body=%s", len(message.ToolCalls), e.ResponseBody)
	}
	tool := message.ToolCalls[0]
	if tool.ID != "call_1" || tool.Type != "function" || tool.Function.Name != "lookup" || tool.Function.Arguments != `{"q":"weather"}` {
		t.Fatalf("unexpected tool call: %+v", tool)
	}
}

func TestOpenAIResponsesStreamAssemblesCompletedResponse(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`event: response.created` + "\n" + `data: {"type":"response.created","response":{"id":"resp-1","object":"response","status":"in_progress","model":"gpt-5.4-mini","output":[]}}` + "\n\n",
		`event: response.output_text.delta` + "\n" + `data: {"type":"response.output_text.delta","item_id":"msg-1","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n",
		`event: response.completed` + "\n" + `data: {"type":"response.completed","response":{"id":"resp-1","object":"response","status":"completed","model":"gpt-5.4-mini","output":[{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/responses")

	var captured struct {
		ID     string `json:"id"`
		Object string `json:"object"`
		Status string `json:"status"`
		Model  string `json:"model"`
		Output []struct {
			Type    string `json:"type"`
			Role    string `json:"role"`
			Content []struct {
				Type string `json:"type"`
				Text string `json:"text"`
			} `json:"content"`
		} `json:"output"`
		Usage struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
			TotalTokens  int `json:"total_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("responses stream should be completed response JSON: %v; body=%q", err, e.ResponseBody)
	}
	if captured.ID != "resp-1" || captured.Object != "response" || captured.Status != "completed" || captured.Model != "gpt-5.4-mini" {
		t.Fatalf("unexpected response metadata: %+v", captured)
	}
	if len(captured.Output) != 1 || len(captured.Output[0].Content) != 1 || captured.Output[0].Content[0].Text != "hello" {
		t.Fatalf("unexpected response output: %+v", captured.Output)
	}
	if captured.Usage.TotalTokens != 4 {
		t.Fatalf("usage not preserved: %+v", captured.Usage)
	}
}

func TestOpenAICompletionsStreamAssemblesText(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"hello","finish_reason":null}]}` + "\n\n",
		`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":" world","finish_reason":null}]}` + "\n\n",
		`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"","finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":2,"total_tokens":4}}` + "\n\n",
		"data: [DONE]\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/completions")

	var captured struct {
		ID      string `json:"id"`
		Object  string `json:"object"`
		Created int64  `json:"created"`
		Model   string `json:"model"`
		Choices []struct {
			Index        int    `json:"index"`
			Text         string `json:"text"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage struct {
			TotalTokens int `json:"total_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("completions stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
	}
	if captured.ID != "cmpl-1" || captured.Object != "text_completion" || captured.Model != "gpt-5.4-mini" || captured.Created != 1779335544 {
		t.Fatalf("unexpected metadata: %+v", captured)
	}
	if len(captured.Choices) != 1 || captured.Choices[0].Text != "hello world" || captured.Choices[0].FinishReason != "stop" {
		t.Fatalf("unexpected choices: %+v", captured.Choices)
	}
	if captured.Usage.TotalTokens != 4 {
		t.Fatalf("usage not preserved: %+v", captured.Usage)
	}
}

func TestAnthropicSSEAssemblesNativeMessageJSON(t *testing.T) {
	e := requestStreamEntry(t, fakeAnthropicStreamUpstream(), "/v1/messages")

	var captured struct {
		ID      string `json:"id"`
		Type    string `json:"type"`
		Role    string `json:"role"`
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
		StopReason string `json:"stop_reason"`
		Usage      struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("anthropic stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
	}
	if captured.ID != "msg-1" || captured.Type != "message" || captured.Role != "assistant" {
		t.Fatalf("unexpected anthropic message metadata: %+v", captured)
	}
	if len(captured.Content) != 1 || captured.Content[0].Type != "text" || captured.Content[0].Text != "你好" {
		t.Fatalf("unexpected anthropic content: %+v", captured.Content)
	}
	if captured.StopReason != "end_turn" {
		t.Errorf("stop_reason=%q, want end_turn", captured.StopReason)
	}
	if captured.Usage.InputTokens != 10 || captured.Usage.OutputTokens != 3 {
		t.Errorf("usage=%+v, want input=10 output=3", captured.Usage)
	}
}

func TestAnthropicSSEAssemblesToolUseContent(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-tool","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
		`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}` + "\n\n",
		`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}` + "\n\n",
		`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"weather\"}"}}` + "\n\n",
		`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":8}}` + "\n\n",
		`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1/messages")

	var captured struct {
		Content []struct {
			Type  string         `json:"type"`
			ID    string         `json:"id"`
			Name  string         `json:"name"`
			Input map[string]any `json:"input"`
		} `json:"content"`
		StopReason string `json:"stop_reason"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("anthropic tool stream should be JSON: %v; body=%q", err, e.ResponseBody)
	}
	if len(captured.Content) != 1 {
		t.Fatalf("content length=%d, want 1; body=%s", len(captured.Content), e.ResponseBody)
	}
	tool := captured.Content[0]
	if tool.Type != "tool_use" || tool.ID != "toolu_1" || tool.Name != "lookup" || tool.Input["q"] != "weather" {
		t.Fatalf("unexpected tool content: %+v", tool)
	}
	if captured.StopReason != "tool_use" {
		t.Fatalf("stop_reason=%q, want tool_use", captured.StopReason)
	}
}

func TestGeminiSSEAssemblesNativeGenerateContentJSON(t *testing.T) {
	e := requestStreamEntry(t, fakeGeminiStreamUpstream(), "/v1beta/models/gemini-1.5-pro:generateContent")

	var captured struct {
		Candidates []struct {
			Content struct {
				Role  string `json:"role"`
				Parts []struct {
					Text string `json:"text"`
				} `json:"parts"`
			} `json:"content"`
			FinishReason string `json:"finishReason"`
			Index        int    `json:"index"`
		} `json:"candidates"`
		UsageMetadata struct {
			PromptTokenCount     int `json:"promptTokenCount"`
			CandidatesTokenCount int `json:"candidatesTokenCount"`
			TotalTokenCount      int `json:"totalTokenCount"`
		} `json:"usageMetadata"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("gemini stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
	}
	if len(captured.Candidates) != 1 {
		t.Fatalf("candidates length=%d, want 1", len(captured.Candidates))
	}
	candidate := captured.Candidates[0]
	if candidate.Content.Role != "model" || len(candidate.Content.Parts) != 1 || candidate.Content.Parts[0].Text != "你好" {
		t.Fatalf("unexpected gemini content: %+v", candidate.Content)
	}
	if candidate.FinishReason != "STOP" {
		t.Errorf("finishReason=%q, want STOP", candidate.FinishReason)
	}
	if captured.UsageMetadata.TotalTokenCount != 4 {
		t.Errorf("usageMetadata=%+v, want totalTokenCount=4", captured.UsageMetadata)
	}
}

func TestGeminiStreamPreservesSafetyRatingsAndUsageMetadata(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"finishReason":null,"index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":8,"thoughtsTokenCount":10}}` + "\n\n",
		`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":2,"totalTokenCount":20,"thoughtsTokenCount":10}}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")

	var captured struct {
		Candidates []struct {
			Content struct {
				Parts []struct {
					Text string `json:"text"`
				} `json:"parts"`
			} `json:"content"`
			SafetyRatings []struct {
				Category    string `json:"category"`
				Probability string `json:"probability"`
			} `json:"safetyRatings"`
		} `json:"candidates"`
		UsageMetadata struct {
			PromptTokenCount        int `json:"promptTokenCount"`
			ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
			CandidatesTokenCount    int `json:"candidatesTokenCount"`
			TotalTokenCount         int `json:"totalTokenCount"`
			ThoughtsTokenCount      int `json:"thoughtsTokenCount"`
		} `json:"usageMetadata"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("gemini stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
	}
	if len(captured.Candidates) != 1 || len(captured.Candidates[0].SafetyRatings) != 1 {
		t.Fatalf("expected safetyRatings to be preserved, got %+v", captured.Candidates)
	}
	if captured.Candidates[0].SafetyRatings[0].Category != "HARM_CATEGORY_HARASSMENT" {
		t.Fatalf("unexpected safetyRatings: %+v", captured.Candidates[0].SafetyRatings)
	}
	if captured.UsageMetadata.ToolUsePromptTokenCount != 0 || captured.UsageMetadata.ThoughtsTokenCount != 10 || captured.UsageMetadata.TotalTokenCount != 20 {
		t.Fatalf("usageMetadata fields not preserved: %+v", captured.UsageMetadata)
	}
}

func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")

	var captured struct {
		Candidates []struct {
			Content struct {
				Parts []struct {
					FunctionCall struct {
						Name string         `json:"name"`
						Args map[string]any `json:"args"`
					} `json:"functionCall"`
				} `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("gemini functionCall stream should be JSON: %v; body=%q", err, e.ResponseBody)
	}
	call := captured.Candidates[0].Content.Parts[0].FunctionCall
	if call.Name != "lookup" || call.Args["q"] != "weather" {
		t.Fatalf("functionCall not preserved: %+v; body=%s", call, e.ResponseBody)
	}
}

func TestGeminiStreamPreservesMultiplePartsByIndex(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"candidates":[{"content":{"parts":[{"text":"hello "},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"world"}],"role":"model"},"index":0}]}` + "\n\n",
		`data: {"candidates":[{"content":{"parts":[{"text":"again"},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"!"}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")

	var captured struct {
		Candidates []struct {
			Content struct {
				Parts []map[string]any `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
	}
	parts := captured.Candidates[0].Content.Parts
	if len(parts) != 3 || parts[0]["text"] != "hello again" || parts[2]["text"] != "world!" {
		t.Fatalf("multiple Gemini parts not preserved: %+v", parts)
	}
	if _, ok := parts[1]["functionCall"]; !ok {
		t.Fatalf("middle functionCall part missing: %+v", parts)
	}
}

func TestGeminiStreamAppendsDifferentPartKindsAcrossChunks(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"candidates":[{"content":{"parts":[{"text":"answer"}],"role":"model"},"index":0}]}` + "\n\n",
		`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
	})
	e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")

	var captured struct {
		Candidates []struct {
			Content struct {
				Parts []map[string]any `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
	}
	if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
		t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
	}
	parts := captured.Candidates[0].Content.Parts
	if len(parts) != 2 || parts[0]["text"] != "answer" {
		t.Fatalf("different Gemini part kinds were merged: %+v", parts)
	}
	if _, ok := parts[1]["functionCall"]; !ok {
		t.Fatalf("functionCall part missing: %+v", parts)
	}
}

func TestUnknownSSEKeepsRawBody(t *testing.T) {
	e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")

	if string(e.ResponseBody) != "event: custom\ndata: not-json\n\n" {
		t.Fatalf("unknown stream should keep raw body, got %q", e.ResponseBody)
	}
}

func TestTruncatedSSEKeepsCapturedRawBody(t *testing.T) {
	upstream := fakeOpenAIStreamUpstream()
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 520)

	proxySrv := httptest.NewServer(h)
	defer proxySrv.Close()

	resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", nil)
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()

	deadline := time.Now().Add(2 * time.Second)
	for sub.Len() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if sub.Len() == 0 {
		t.Fatal("no log entry captured")
	}
	e := sub.Entry(0)

	if !e.ResponseTruncated {
		t.Fatal("response should be marked truncated")
	}
	if !strings.HasPrefix(string(e.ResponseBody), "data: ") {
		t.Fatalf("truncated stream should keep captured raw body, got %q", e.ResponseBody)
	}
	if json.Valid(e.ResponseBody) {
		t.Fatalf("truncated stream should not be assembled as JSON, got %q", e.ResponseBody)
	}
}

func TestMultimodalStreamRequestBodyIsCaptured(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
		`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
		"data: [DONE]\n\n",
	})
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 1024*1024)

	proxySrv := httptest.NewServer(h)
	defer proxySrv.Close()

	reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
	resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(reqBody))
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()

	deadline := time.Now().Add(2 * time.Second)
	for sub.Len() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if sub.Len() == 0 {
		t.Fatal("no log entry captured")
	}
	e := sub.Entry(0)

	if e.RequestTruncated {
		t.Fatal("multimodal request should not be truncated")
	}
	requestText := string(e.RequestBody)
	if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
		t.Fatalf("request_body should contain image input, got %q", requestText)
	}
	if !e.IsStream {
		t.Fatal("response should still be marked stream")
	}
}

func TestChunkedMultimodalRequestBodyIsCaptured(t *testing.T) {
	upstream := newSSEUpstream([]string{
		`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}` + "\n\n",
		"data: [DONE]\n\n",
	})
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 1024*1024)

	proxySrv := httptest.NewServer(h)
	defer proxySrv.Close()

	reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
	req, err := http.NewRequest(http.MethodPost, proxySrv.URL+"/v1/chat/completions", strings.NewReader(reqBody))
	if err != nil {
		t.Fatal(err)
	}
	req.ContentLength = -1
	req.Header.Set("Content-Type", "application/json")
	req.TransferEncoding = []string{"chunked"}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()

	deadline := time.Now().Add(2 * time.Second)
	for sub.Len() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if sub.Len() == 0 {
		t.Fatal("no log entry captured")
	}
	e := sub.Entry(0)

	requestText := string(e.RequestBody)
	if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
		t.Fatalf("chunked request_body should contain image input, got %q", requestText)
	}
}

func requestStreamEntry(t *testing.T, upstream *httptest.Server, path string) *logger.LogEntry {
	t.Helper()
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)

	sub := &captureSubmitter{}
	filter, err := config.NewFilter(config.FilterDisabled, nil)
	if err != nil {
		t.Fatal(err)
	}
	h := proxy.New(u, filter, sub, 1024*1024)

	proxySrv := httptest.NewServer(h)
	defer proxySrv.Close()

	resp, err := http.Post(proxySrv.URL+path, "application/json", nil)
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()

	deadline := time.Now().Add(2 * time.Second)
	for sub.Len() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if sub.Len() == 0 {
		t.Fatal("no log entry captured")
	}
	return sub.Entry(0)
}