package llm

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// sse formats events the way the API does.
func sse(events ...string) string {
	var b strings.Builder
	for _, e := range events {
		b.WriteString("event: x\ndata: " + e + "\n\n")
	}
	return b.String()
}

func streamServer(t *testing.T, payload string) *Client {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("accept") != "text/event-stream" {
			t.Errorf("accept = %q", r.Header.Get("accept"))
		}
		var req map[string]any
		body, _ := io.ReadAll(r.Body)
		json.Unmarshal(body, &req)
		if req["stream"] != true {
			t.Errorf("stream flag not set in body: %s", body)
		}
		w.Header().Set("content-type", "text/event-stream")
		io.WriteString(w, payload)
	}))
	t.Cleanup(srv.Close)
	c := New(WithBaseURL(srv.URL), WithAPIKey("k"))
	c.RetryDelay = time.Millisecond
	return c
}

// The fiddly part: tool arguments arrive as fragments that are individually
// invalid JSON and only parseable once concatenated.
func TestStreamAccumulatesToolInputFragments(t *testing.T) {
	c := streamServer(t, sse(
		`{"type":"message_start","message":{"id":"msg_1","model":"claude-opus-5","usage":{"input_tokens":100,"cache_read_input_tokens":40}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Let me "}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"look."}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tu_1","name":"edit_file","input":{}}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"path\":"}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"main.go\",\"old_str\""}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":":\"a\",\"new_str\":\"b\"}"}}`,
		`{"type":"content_block_stop","index":1}`,
		`{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":55}}`,
		`{"type":"message_stop"}`,
	))

	var textSeen strings.Builder
	var toolStarts int
	resp, err := c.StreamMessage(context.Background(), Request{}, func(ev StreamEvent) {
		switch ev.Kind {
		case StreamText:
			textSeen.WriteString(ev.Text)
		case StreamToolStart:
			toolStarts++
			if ev.ToolName != "edit_file" || ev.ToolID != "tu_1" {
				t.Errorf("tool start = %+v", ev)
			}
		}
	})
	if err != nil {
		t.Fatalf("StreamMessage: %v", err)
	}
	if textSeen.String() != "Let me look." {
		t.Errorf("streamed text = %q", textSeen.String())
	}
	if toolStarts != 1 {
		t.Errorf("%d tool starts, want 1", toolStarts)
	}
	if resp.StopReason != StopToolUse {
		t.Errorf("stop_reason = %q", resp.StopReason)
	}
	if got := resp.Text(); got != "Let me look." {
		t.Errorf("assembled text = %q", got)
	}

	uses := resp.ToolUses()
	if len(uses) != 1 {
		t.Fatalf("%d tool uses, want 1", len(uses))
	}
	var args struct{ Path, OldStr, NewStr string }
	if err := json.Unmarshal(uses[0].Input, &struct {
		Path   *string `json:"path"`
		OldStr *string `json:"old_str"`
		NewStr *string `json:"new_str"`
	}{&args.Path, &args.OldStr, &args.NewStr}); err != nil {
		t.Fatalf("reassembled input is not valid JSON: %v (%s)", err, uses[0].Input)
	}
	if args.Path != "main.go" || args.OldStr != "a" || args.NewStr != "b" {
		t.Errorf("reassembled args = %+v", args)
	}
	if resp.Usage.InputTokens != 100 || resp.Usage.CacheReadInputTokens != 40 || resp.Usage.OutputTokens != 55 {
		t.Errorf("usage = %+v: input counts come from message_start, output from message_delta", resp.Usage)
	}
}

// A hosted web search streams like a local tool call — server_tool_use with
// "input":{} at content_block_start, the query as input_json_delta fragments —
// and then the API drops the whole web_search_tool_result in as one block.
// The conversation is re-sent after pause_turn, so what matters is that the
// re-marshalled message carries the query and the result exactly as received.
func TestStreamServerToolUseRoundTripsItsQuery(t *testing.T) {
	const result = `{"type":"web_search_tool_result","tool_use_id":"srvtoolu_1","content":[{"type":"web_search_result","title":"Go 1.23","url":"https://go.dev/doc/go1.23","encrypted_content":"abc==","page_age":null}]}`
	c := streamServer(t, sse(
		`{"type":"message_start","message":{"id":"m","usage":{"input_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"server_tool_use","id":"srvtoolu_1","name":"web_search","input":{}}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"query\":"}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"go 1.23 release\"}"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"content_block_start","index":1,"content_block":`+result+`}`,
		`{"type":"content_block_stop","index":1}`,
		`{"type":"message_delta","delta":{"stop_reason":"pause_turn"},"usage":{"output_tokens":1}}`,
	))
	resp, err := c.StreamMessage(context.Background(), Request{}, nil)
	if err != nil {
		t.Fatalf("StreamMessage: %v", err)
	}
	if resp.StopReason != StopPauseTurn {
		t.Errorf("stop_reason = %q, want pause_turn", resp.StopReason)
	}
	if len(resp.ToolUses()) != 0 {
		t.Errorf("ToolUses() = %+v, want none: a hosted call is not for the local registry", resp.ToolUses())
	}
	server := resp.ServerToolUses()
	if len(server) != 1 || string(server[0].Input) != `{"query":"go 1.23 release"}` {
		t.Fatalf("ServerToolUses() = %+v, want the reassembled query", server)
	}

	out, err := json.Marshal(resp.Message())
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if !strings.Contains(string(out), `"input":{"query":"go 1.23 release"}`) {
		t.Errorf("re-sent server_tool_use lost its query — the content_block_start snapshot was sent instead:\n%s", out)
	}
	if !strings.Contains(string(out), result) {
		t.Errorf("web_search_tool_result must go back verbatim:\n%s", out)
	}
}

func TestStreamPreservesThinkingSignature(t *testing.T) {
	c := streamServer(t, sse(
		`{"type":"message_start","message":{"id":"m","usage":{"input_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig-"}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"abc"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}`,
	))
	resp, err := c.StreamMessage(context.Background(), Request{}, nil)
	if err != nil {
		t.Fatalf("StreamMessage: %v", err)
	}
	if len(resp.Content) != 1 || resp.Content[0].Signature != "sig-abc" {
		t.Errorf("signature = %q, want the fragments joined", resp.Content[0].Signature)
	}
	out, _ := json.Marshal(resp.Content[0])
	if !strings.Contains(string(out), `"signature":"sig-abc"`) {
		t.Errorf("signature must survive re-marshalling: %s", out)
	}
}

func TestStreamErrors(t *testing.T) {
	tests := []struct {
		name    string
		payload string
		want    string
	}{
		{
			name: "mid-stream error event",
			payload: sse(
				`{"type":"message_start","message":{"id":"m"}}`,
				`{"type":"error","error":{"type":"overloaded_error","message":"try later"}}`),
			want: "try later",
		},
		{
			name: "truncated tool arguments",
			payload: sse(
				`{"type":"message_start","message":{"id":"m"}}`,
				`{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t","name":"bash"}}`,
				`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"command\":"}}`,
				`{"type":"content_block_stop","index":0}`,
				`{"type":"message_delta","delta":{"stop_reason":"tool_use"}}`),
			want: "invalid JSON arguments",
		},
		{
			name:    "stream ends without a stop reason",
			payload: sse(`{"type":"message_start","message":{"id":"m"}}`),
			want:    "without a stop_reason",
		},
		{
			name:    "malformed event",
			payload: "data: {not json\n\n",
			want:    "malformed stream event",
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := streamServer(t, tc.payload).StreamMessage(context.Background(), Request{}, nil)
			if err == nil || !strings.Contains(err.Error(), tc.want) {
				t.Errorf("error = %v, want it to mention %q", err, tc.want)
			}
		})
	}
}

func TestStreamToolUseWithNoArguments(t *testing.T) {
	c := streamServer(t, sse(
		`{"type":"message_start","message":{"id":"m"}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t","name":"pwd"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"message_delta","delta":{"stop_reason":"tool_use"}}`,
	))
	resp, err := c.StreamMessage(context.Background(), Request{}, nil)
	if err != nil {
		t.Fatalf("StreamMessage: %v", err)
	}
	if got := string(resp.ToolUses()[0].Input); got != "{}" {
		t.Errorf("input = %s, want an empty object", got)
	}
}

func TestStreamHTTPErrorIsAnAPIError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, `{"error":{"type":"invalid_request_error","message":"bad model"}}`)
	}))
	defer srv.Close()

	c := New(WithBaseURL(srv.URL), WithAPIKey("k"))
	_, err := c.StreamMessage(context.Background(), Request{}, nil)
	apiErr, ok := err.(*APIError)
	if !ok {
		t.Fatalf("error type = %T, want *APIError", err)
	}
	if apiErr.StatusCode != 400 || apiErr.Message != "bad model" {
		t.Errorf("APIError = %+v", apiErr)
	}
}