package llm

import (
	"bytes"
	"encoding/json"
	"sort"
	"strings"
	"testing"
)

func TestContentBlockMarshalsOnlyItsOwnFields(t *testing.T) {
	tests := []struct {
		name  string
		block ContentBlock
		want  string
	}{
		{
			name:  "text carries no stray thinking key",
			block: TextBlock("hello"),
			want:  `{"text":"hello","type":"text"}`,
		},
		{
			name:  "tool_use keeps input as raw JSON",
			block: ContentBlock{Type: BlockToolUse, ID: "tu_1", Name: "read_file", Input: json.RawMessage(`{"path":"main.go"}`)},
			want:  `{"id":"tu_1","input":{"path":"main.go"},"name":"read_file","type":"tool_use"}`,
		},
		{
			name:  "tool_use with no input sends an empty object",
			block: ContentBlock{Type: BlockToolUse, ID: "tu_2", Name: "pwd"},
			want:  `{"id":"tu_2","input":{},"name":"pwd","type":"tool_use"}`,
		},
		{
			name:  "successful tool_result omits is_error",
			block: ToolResultBlock("tu_1", "ok", false),
			want:  `{"content":"ok","tool_use_id":"tu_1","type":"tool_result"}`,
		},
		{
			name:  "failed tool_result sets is_error",
			block: ToolResultBlock("tu_1", "no such file", true),
			want:  `{"content":"no such file","is_error":true,"tool_use_id":"tu_1","type":"tool_result"}`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got, err := json.Marshal(tc.block)
			if err != nil {
				t.Fatalf("marshal: %v", err)
			}
			if string(got) != tc.want {
				t.Errorf("got  %s\nwant %s", got, tc.want)
			}
		})
	}
}

// With display omitted (the default on Opus 5) a thinking block has empty text
// but a signature that must go back unchanged, so the field cannot be omitempty.
func TestThinkingBlockRoundTripsEmptyTextAndSignature(t *testing.T) {
	const wire = `{"type":"thinking","thinking":"","signature":"sig-abc123"}`
	var b ContentBlock
	if err := json.Unmarshal([]byte(wire), &b); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	out, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if !strings.Contains(string(out), `"thinking":""`) {
		t.Errorf("empty thinking text must survive the round trip, got %s", out)
	}
	if !strings.Contains(string(out), `"signature":"sig-abc123"`) {
		t.Errorf("signature must survive the round trip, got %s", out)
	}
}

func TestUnknownBlockTypeRoundTripsVerbatim(t *testing.T) {
	const wire = `{"type":"server_tool_use_future","id":"x","weird":{"nested":[1,2]}}`
	var b ContentBlock
	if err := json.Unmarshal([]byte(wire), &b); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	out, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if string(out) != wire {
		t.Errorf("unmodelled block should be echoed verbatim\ngot  %s\nwant %s", out, wire)
	}
}

// A hosted tool is declared by type, not by schema. The API rejects a tools
// entry that carries "input_schema":null, which is what a non-omitempty field
// would send for it — so this asserts on the set of keys, not on their values.
func TestWebSearchToolDefHasNoLocalSchema(t *testing.T) {
	tests := []struct {
		name     string
		def      ToolDef
		wantKeys []string
	}{
		{
			name:     "hosted web search with a cap",
			def:      WebSearchTool(5),
			wantKeys: []string{"max_uses", "name", "type"},
		},
		{
			name:     "hosted web search with no cap omits max_uses",
			def:      WebSearchTool(0),
			wantKeys: []string{"name", "type"},
		},
		{
			name:     "an ordinary local tool still sends its schema",
			def:      ToolDef{Name: "read_file", Description: "Read a file", InputSchema: json.RawMessage(`{"type":"object"}`)},
			wantKeys: []string{"description", "input_schema", "name"},
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			raw, err := json.Marshal(tc.def)
			if err != nil {
				t.Fatalf("marshal: %v", err)
			}
			var obj map[string]json.RawMessage
			if err := json.Unmarshal(raw, &obj); err != nil {
				t.Fatalf("unmarshal: %v", err)
			}
			var keys []string
			for k := range obj {
				keys = append(keys, k)
			}
			sort.Strings(keys)
			if strings.Join(keys, ",") != strings.Join(tc.wantKeys, ",") {
				t.Errorf("keys = %v, want %v\n%s", keys, tc.wantKeys, raw)
			}
		})
	}

	got, _ := json.Marshal(WebSearchTool(5))
	if want := `{"type":"web_search_20250305","name":"web_search","max_uses":5}`; string(got) != want {
		t.Errorf("WebSearchTool(5) = %s, want %s", got, want)
	}
}

// A streamed server_tool_use starts with "input":{} and gets its query as
// input_json_delta fragments. Re-sending the snapshot from content_block_start
// would drop the query, so the block must be marshalled from its fields.
func TestServerToolUseMarshalsAccumulatedInput(t *testing.T) {
	var b ContentBlock
	if err := json.Unmarshal([]byte(`{"type":"server_tool_use","id":"srvtoolu_1","name":"web_search","input":{}}`), &b); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	b.Input = json.RawMessage(`{"query":"go 1.23 release notes"}`)
	out, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	want := `{"id":"srvtoolu_1","input":{"query":"go 1.23 release notes"},"name":"web_search","type":"server_tool_use"}`
	if string(out) != want {
		t.Errorf("got  %s\nwant %s", out, want)
	}

	// With no input at all the API still wants an object, as for tool_use.
	empty := ContentBlock{Type: BlockServerToolUse, ID: "srvtoolu_2", Name: WebSearchToolName}
	out, _ = json.Marshal(empty)
	if !strings.Contains(string(out), `"input":{}`) {
		t.Errorf("empty input must become {}, got %s", out)
	}
}

// web_search_tool_result carries content as an array (or an error object),
// where the struct models content as a string. It must still parse — a block
// that does not parse takes the whole response down with it — and go back
// byte for byte, because the API supplied it and Codey has nothing to add.
func TestWebSearchToolResultRoundTripsVerbatim(t *testing.T) {
	tests := []struct {
		name string
		wire string
	}{
		{
			name: "results array",
			wire: `{"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":"2 days ago"}]}`,
		},
		{
			name: "error object",
			wire: `{"type":"web_search_tool_result","tool_use_id":"srvtoolu_1","content":{"type":"web_search_tool_result_error","error_code":"max_uses_exceeded"}}`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var b ContentBlock
			if err := json.Unmarshal([]byte(tc.wire), &b); err != nil {
				t.Fatalf("unmarshal: %v", err)
			}
			if b.Type != BlockWebSearchToolResult || b.ToolUseID != "srvtoolu_1" {
				t.Errorf("type/tool_use_id = %q/%q, want the modelled fields filled in", b.Type, b.ToolUseID)
			}
			if string(b.Raw()) != tc.wire {
				t.Errorf("Raw() = %s, want the wire bytes", b.Raw())
			}
			out, err := json.Marshal(b)
			if err != nil {
				t.Fatalf("marshal: %v", err)
			}
			if string(out) != tc.wire {
				t.Errorf("got  %s\nwant %s", out, tc.wire)
			}

			// The loop puts a cache breakpoint on the last block of the last
			// message, and after pause_turn that block is very often this one.
			b.CacheControl = Ephemeral()
			out, err = json.Marshal(b)
			if err != nil {
				t.Fatalf("marshal with cache_control: %v", err)
			}
			var obj map[string]json.RawMessage
			if err := json.Unmarshal(out, &obj); err != nil {
				t.Fatalf("re-marshalled block is not an object: %v", err)
			}
			if string(obj["cache_control"]) != `{"type":"ephemeral"}` {
				t.Errorf("cache_control = %s, want it attached to the raw block", obj["cache_control"])
			}
			var orig map[string]json.RawMessage
			json.Unmarshal([]byte(tc.wire), &orig)
			for k, v := range orig {
				if string(obj[k]) != string(v) {
					t.Errorf("%s = %s after attaching cache_control, want %s", k, obj[k], v)
				}
			}
		})
	}
}

// A hosted block must never reach the local registry: ToolUses filters on
// tool_use alone, and ServerToolUses is the separate view for the observer.
func TestServerToolUsesAreNotToolUses(t *testing.T) {
	const wire = `{
	  "id":"msg_1","type":"message","role":"assistant","model":"claude-opus-5",
	  "content":[
	    {"type":"server_tool_use","id":"srvtoolu_1","name":"web_search","input":{"query":"a"}},
	    {"type":"web_search_tool_result","tool_use_id":"srvtoolu_1","content":[]},
	    {"type":"tool_use","id":"tu_1","name":"read_file","input":{"path":"a.go"}},
	    {"type":"server_tool_use","id":"srvtoolu_2","name":"web_search","input":{"query":"b"}}
	  ],
	  "stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}
	}`
	var r Response
	if err := json.Unmarshal([]byte(wire), &r); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	uses := r.ToolUses()
	if len(uses) != 1 || uses[0].ID != "tu_1" {
		t.Errorf("ToolUses() = %+v, want only the local read_file call", uses)
	}
	server := r.ServerToolUses()
	if len(server) != 2 || server[0].ID != "srvtoolu_1" || server[1].ID != "srvtoolu_2" {
		t.Errorf("ServerToolUses() = %+v, want both hosted calls in order", server)
	}
}

// A tool_result whose content is blocks — a caption and an image — goes to the
// API as an array, and comes back from a saved session the same way. Text-only
// readers still see the caption in Content.
func TestToolResultWithBlocksRoundTrips(t *testing.T) {
	b := ToolResultBlocks("tu_1", []ContentBlock{
		TextBlock("PNG image, 2×2, 70 bytes."),
		ImageBlock("image/png", "AAAA"),
	})
	if b.Content != "PNG image, 2×2, 70 bytes." {
		t.Errorf("Content = %q, want the caption so text-only readers have something", b.Content)
	}

	wire, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	const want = `{"content":[{"text":"PNG image, 2×2, 70 bytes.","type":"text"},{"source":{"type":"base64","media_type":"image/png","data":"AAAA"},"type":"image"}],"tool_use_id":"tu_1","type":"tool_result"}`
	if string(wire) != want {
		t.Errorf("wire:\n got %s\nwant %s", wire, want)
	}

	var back ContentBlock
	if err := json.Unmarshal(wire, &back); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if len(back.Blocks) != 2 || back.Blocks[1].Type != BlockImage || back.Blocks[1].Source == nil || back.Blocks[1].Source.Data != "AAAA" {
		t.Errorf("blocks did not survive the round trip: %+v", back.Blocks)
	}
	if back.Content != b.Content {
		t.Errorf("Content after round trip = %q, want %q", back.Content, b.Content)
	}
	again, err := json.Marshal(back)
	if err != nil {
		t.Fatalf("re-marshal: %v", err)
	}
	if string(again) != want {
		t.Errorf("second marshal differs:\n got %s\nwant %s", again, want)
	}
}

func TestToolResultWithNoBlocksStillMarshalsAString(t *testing.T) {
	wire, err := json.Marshal(ToolResultBlocks("tu_1", nil))
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if string(wire) != `{"content":"(no output)","tool_use_id":"tu_1","type":"tool_result"}` {
		t.Errorf("got %s", wire)
	}
}

func TestEmptyToolResultGetsPlaceholder(t *testing.T) {
	if got := ToolResultBlock("tu_1", "", false).Content; got != "(no output)" {
		t.Errorf("content = %q, want a placeholder: the API rejects empty tool_result content", got)
	}
}

func TestResponseAccessors(t *testing.T) {
	const wire = `{
	  "id":"msg_1","type":"message","role":"assistant","model":"claude-opus-5",
	  "content":[
	    {"type":"thinking","thinking":"","signature":"s"},
	    {"type":"text","text":"Let me look. "},
	    {"type":"tool_use","id":"tu_1","name":"read_file","input":{"path":"a.go"}},
	    {"type":"text","text":"One moment."}
	  ],
	  "stop_reason":"tool_use","stop_details":null,
	  "usage":{"input_tokens":10,"output_tokens":20,"cache_read_input_tokens":5,"cache_creation_input_tokens":1}
	}`
	var r Response
	if err := json.Unmarshal([]byte(wire), &r); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if got := r.Text(); got != "Let me look. One moment." {
		t.Errorf("Text() = %q, want the text blocks concatenated", got)
	}
	uses := r.ToolUses()
	if len(uses) != 1 || uses[0].ID != "tu_1" || uses[0].Name != "read_file" {
		t.Errorf("ToolUses() = %+v, want the single read_file call", uses)
	}
	if r.Usage.CacheReadInputTokens != 5 {
		t.Errorf("cache_read_input_tokens = %d, want 5", r.Usage.CacheReadInputTokens)
	}
	if msg := r.Message(); msg.Role != RoleAssistant || len(msg.Content) != 4 {
		t.Errorf("Message() = %+v, want all 4 blocks under the assistant role", msg)
	}
}

// The effort levels are a closed set, and both the -effort flag and the TUI's
// /effort command answer out of it rather than keeping their own copy.
func TestValidEffort(t *testing.T) {
	for _, level := range EffortLevels {
		if !ValidEffort(level) {
			t.Errorf("ValidEffort(%q) = false for a level we advertise", level)
		}
	}
	for _, level := range []string{"", "HIGH", "higher", "none"} {
		if ValidEffort(level) {
			t.Errorf("ValidEffort(%q) = true, want false", level)
		}
	}
}

// An image block carries a nested source object, and the API is particular
// about its shape: type, media_type and data, and nothing else alongside.
func TestImageBlockMarshalsAsASource(t *testing.T) {
	b, err := json.Marshal(ImageBlock("image/png", "aGVsbG8="))
	if err != nil {
		t.Fatal(err)
	}
	var got map[string]any
	if err := json.Unmarshal(b, &got); err != nil {
		t.Fatal(err)
	}
	if got["type"] != "image" {
		t.Errorf("type = %v", got["type"])
	}
	src, ok := got["source"].(map[string]any)
	if !ok {
		t.Fatalf("source = %v", got["source"])
	}
	if src["type"] != "base64" || src["media_type"] != "image/png" || src["data"] != "aGVsbG8=" {
		t.Errorf("source = %v", src)
	}
	// A stray "text" key on an image block is the kind of thing the API
	// rejects outright.
	if _, stray := got["text"]; stray || len(got) != 2 {
		t.Errorf("extra keys on an image block: %v", got)
	}
}

// A picture has to survive being written out and read back, because that is
// what happens to every earlier turn on every subsequent request.
func TestImageBlockRoundTrips(t *testing.T) {
	msg := Message{Role: RoleUser, Content: []ContentBlock{
		ImageBlock("image/jpeg", "Zm9v"),
		TextBlock("what is this?"),
	}}
	b, err := json.Marshal(msg)
	if err != nil {
		t.Fatal(err)
	}
	var back Message
	if err := json.Unmarshal(b, &back); err != nil {
		t.Fatal(err)
	}
	if len(back.Content) != 2 || back.Content[0].Source == nil {
		t.Fatalf("content = %+v", back.Content)
	}
	if got := back.Content[0].Source; got.MediaType != "image/jpeg" || got.Data != "Zm9v" {
		t.Errorf("source = %+v", got)
	}
	again, err := json.Marshal(back)
	if err != nil || !bytes.Equal(b, again) {
		t.Errorf("second trip differs:\n%s\n%s\n%v", b, again, err)
	}
}