package ai

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

func TestRepairChatCompletionBody_preservesToolCalls(t *testing.T) {
	body := []byte(`{
		"model":"gpt-4o",
		"messages":[
			{"role":"user","content":"list mail"},
			{"role":"assistant","content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"mail_list","arguments":"{\"limit\":3}"}}]},
			{"role":"tool","tool_call_id":"call_1","content":"{\"items\":[]}"}
		]
	}`)

	repaired, err := repairChatCompletionBody(body)
	if err != nil {
		t.Fatalf("repairChatCompletionBody() error = %v", err)
	}

	var parsed map[string]any
	if err := json.Unmarshal(repaired, &parsed); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	messages := parsed["messages"].([]any)
	assistant := messages[1].(map[string]any)
	toolCalls, ok := assistant["tool_calls"].([]any)
	if !ok || len(toolCalls) != 1 {
		t.Fatalf("expected assistant tool_calls preserved, got %#v", assistant)
	}
}

func TestRepairChatCompletionBody_insertsMissingAssistantToolCalls(t *testing.T) {
	body := []byte(`{
		"model":"gpt-4o",
		"messages":[
			{"role":"system","content":"prompt"},
			{"role":"user","content":"list mail"},
			{"role":"tool","tool_call_id":"call_abc","content":"{\"items\":[]}"}
		]
	}`)

	repaired, err := repairChatCompletionBody(body)
	if err != nil {
		t.Fatalf("repairChatCompletionBody() error = %v", err)
	}

	var parsed map[string]any
	if err := json.Unmarshal(repaired, &parsed); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	messages := parsed["messages"].([]any)
	if len(messages) != 4 {
		t.Fatalf("messages len = %d, want 4", len(messages))
	}
	assistant := messages[2].(map[string]any)
	if messageRole(assistant) != "assistant" {
		t.Fatalf("messages[2].role = %q, want assistant", messageRole(assistant))
	}
	toolCalls, ok := assistant["tool_calls"].([]any)
	if !ok || len(toolCalls) != 1 {
		t.Fatalf("expected inserted tool_calls, got %#v", assistant)
	}
	call := toolCalls[0].(map[string]any)
	if call["id"] != "call_abc" {
		t.Fatalf("tool_call id = %v, want call_abc", call["id"])
	}
}

func TestRepairChatCompletionBody_stripsNothingWhenAlreadyValid(t *testing.T) {
	body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
	repaired, err := repairChatCompletionBody(body)
	if err != nil {
		t.Fatalf("repairChatCompletionBody() error = %v", err)
	}
	if string(repaired) != string(body) {
		t.Fatalf("expected unchanged body, got %s", repaired)
	}
}

func TestRepairToolMessages_doesNotDuplicateAssistant(t *testing.T) {
	messages := []any{
		map[string]any{"role": "user", "content": "list"},
		map[string]any{
			"role":    "assistant",
			"content": "",
			"tool_calls": []any{
				map[string]any{"id": "call_1", "type": "function"},
			},
		},
		map[string]any{"role": "tool", "tool_call_id": "call_1", "content": "{}"},
	}

	repaired := repairToolMessages(messages)
	if len(repaired) != 3 {
		t.Fatalf("len = %d, want 3", len(repaired))
	}
	if !strings.Contains(string(mustJSON(repaired)), "call_1") {
		t.Fatal("expected original tool_call_id preserved")
	}
}

func mustJSON(v any) []byte {
	b, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return b
}