package canonical

import (
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGetSystemString(t *testing.T) {
	tests := []struct {
		name   string
		system any
		want   string
	}{
		{"string", "hello", "hello"},
		{"nil", nil, ""},
		{"empty string", "", ""},
		{"system blocks", []SystemBlock{{Text: "part1"}, {Text: "part2"}}, "part1\n\npart2"},
		{"single block", []SystemBlock{{Text: "only"}}, "only"},
		{"other type", 123, "123"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := &CanonicalRequest{System: tt.system}
			assert.Equal(t, tt.want, req.GetSystemString())
		})
	}
}

func TestSetSystemString(t *testing.T) {
	req := &CanonicalRequest{}

	req.SetSystemString("hello")
	assert.Equal(t, "hello", req.System)

	req.SetSystemString("")
	assert.Nil(t, req.System)
}

func TestNewTextBlock(t *testing.T) {
	b := NewTextBlock("hello")
	assert.Equal(t, "text", b.Type)
	assert.Equal(t, "hello", b.Text)
}

func TestNewToolUseBlock(t *testing.T) {
	input := json.RawMessage(`{"key":"val"}`)
	b := NewToolUseBlock("id-1", "tool_name", input)
	assert.Equal(t, "tool_use", b.Type)
	assert.Equal(t, "id-1", b.ID)
	assert.Equal(t, "tool_name", b.Name)
	assert.Equal(t, input, b.Input)
}

func TestNewToolResultBlock(t *testing.T) {
	b := NewToolResultBlock("tool-1", "result", false)
	assert.Equal(t, "tool_result", b.Type)
	assert.Equal(t, "tool-1", b.ToolUseID)
	assert.NotNil(t, b.IsError)
	assert.False(t, *b.IsError)
}

func TestNewThinkingBlock(t *testing.T) {
	b := NewThinkingBlock("thought")
	assert.Equal(t, "thinking", b.Type)
	assert.Equal(t, "thought", b.Thinking)
}

func TestNewToolChoice(t *testing.T) {
	assert.Equal(t, &ToolChoice{Type: "auto"}, NewToolChoiceAuto())
	assert.Equal(t, &ToolChoice{Type: "none"}, NewToolChoiceNone())
	assert.Equal(t, &ToolChoice{Type: "any"}, NewToolChoiceAny())
	assert.Equal(t, &ToolChoice{Type: "tool", Name: "fn"}, NewToolChoiceNamed("fn"))
}

func TestCanonicalRequest_RoundTrip(t *testing.T) {
	req := &CanonicalRequest{
		Model:    "gpt-4",
		System:   "system prompt",
		Messages: []CanonicalMessage{{Role: RoleUser, Content: []ContentBlock{NewTextBlock("hi")}}},
		Stream:   true,
	}

	data, err := json.Marshal(req)
	require.NoError(t, err)

	var decoded CanonicalRequest
	require.NoError(t, json.Unmarshal(data, &decoded))
	assert.Equal(t, "gpt-4", decoded.Model)
	assert.Equal(t, "system prompt", decoded.System)
	assert.True(t, decoded.Stream)
}

func TestCanonicalResponse_RoundTrip(t *testing.T) {
	sr := StopReasonEndTurn
	resp := &CanonicalResponse{
		ID:         "resp-1",
		Model:      "gpt-4",
		Content:    []ContentBlock{NewTextBlock("hello")},
		StopReason: &sr,
		Usage:      CanonicalUsage{InputTokens: 10, OutputTokens: 5},
	}

	data, err := json.Marshal(resp)
	require.NoError(t, err)

	var decoded CanonicalResponse
	require.NoError(t, json.Unmarshal(data, &decoded))
	assert.Equal(t, "resp-1", decoded.ID)
	assert.Equal(t, StopReasonEndTurn, *decoded.StopReason)
}
