package main

import (
  "net/http/httptest"
  "strings"
  "testing"
)

func TestFormatPromptToolResolutionAndGrouping(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen/qwen3.6-27b",
    Tools: []Tool{
      {
        Type: "function",
        Function: map[string]interface{}{
          "name":        "calculator",
          "description": "Calculate math expression",
        },
      },
    },
    Messages: []ChatMessage{
      {
        Role:    "user",
        Content: "Calculate 25 * 4 and then tell me what the square root of that is.",
      },
      {
        Role: "assistant",
        ToolCalls: []ToolCall{
          {
            ID:   "call_abc123",
            Type: "function",
            Function: ToolCallFunction{
              Name:      "calculator",
              Arguments: `{"expr":"25*4"}`,
            },
          },
        },
      },
      {
        Role:       "tool",
        ToolCallID: "call_abc123",
        Content:    "100",
      },
    },
  }

  prompt := FormatPrompt(req)

  // Verify tool name was resolved from assistant tool call
  if !strings.Contains(prompt, `"name": "calculator"`) {
    t.Errorf("Expected prompt to contain resolved tool name calculator, got:\n%s", prompt)
  }

  // Verify tool_call_id is preserved
  if !strings.Contains(prompt, `"tool_call_id": "call_abc123"`) {
    t.Errorf("Expected prompt to contain tool_call_id, got:\n%s", prompt)
  }

  // Verify Tool Execution Results block exists
  if !strings.Contains(prompt, "[Tool Execution Results]") {
    t.Errorf("Expected prompt to have [Tool Execution Results], got:\n%s", prompt)
  }

  // Verify Next Steps Directive exists
  if !strings.Contains(prompt, "[Next Steps Directive]") {
    t.Errorf("Expected prompt to have [Next Steps Directive], got:\n%s", prompt)
  }

  // Verify tool result was not placed under [User]
  if strings.Contains(prompt, "[User]\n<tool_response>") {
    t.Errorf("Tool result should not be placed under [User], got:\n%s", prompt)
  }
}

func TestFormatPromptParallelToolResults(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen/qwen3.6-27b",
    Messages: []ChatMessage{
      {
        Role:    "user",
        Content: "What is the weather in Tokyo and Paris?",
      },
      {
        Role: "assistant",
        ToolCalls: []ToolCall{
          {
            ID:   "call_tokyo",
            Type: "function",
            Function: ToolCallFunction{
              Name:      "get_weather",
              Arguments: `{"city":"Tokyo"}`,
            },
          },
          {
            ID:   "call_paris",
            Type: "function",
            Function: ToolCallFunction{
              Name:      "get_weather",
              Arguments: `{"city":"Paris"}`,
            },
          },
        },
      },
      {
        Role:       "tool",
        ToolCallID: "call_tokyo",
        Content:    `{"temp": 18, "condition": "sunny"}`,
      },
      {
        Role:       "tool",
        ToolCallID: "call_paris",
        Content:    `{"temp": 15, "condition": "rainy"}`,
      },
    },
  }

  prompt := FormatPrompt(req)

  // Both tool results must be grouped under [Tool Execution Results]
  execIdx := strings.Index(prompt, "[Tool Execution Results]")
  if execIdx == -1 {
    t.Fatalf("Expected [Tool Execution Results], got:\n%s", prompt)
  }
  execBlock := prompt[execIdx:]

  if !strings.Contains(execBlock, "call_tokyo") {
    t.Errorf("Expected call_tokyo in Tool Execution Results block, got:\n%s", execBlock)
  }
  if !strings.Contains(execBlock, "call_paris") {
    t.Errorf("Expected call_paris in Tool Execution Results block, got:\n%s", execBlock)
  }

  // Neither should be in [Conversation History] as an orphan turn
  historyIdx := strings.Index(prompt, "[Conversation History]")
  if historyIdx != -1 {
    historyBlock := prompt[historyIdx:execIdx]
    if strings.Contains(historyBlock, "call_tokyo") || strings.Contains(historyBlock, "call_paris") {
      t.Errorf("Trailing tool results should not be in history block, got:\n%s", historyBlock)
    }
  }
}

func TestDetectToolCallsFormats(t *testing.T) {
  allowedTools := map[string]bool{
    "search":      true,
    "weather":     true,
    "tool_a":      true,
    "tool_b":      true,
    "direct_tool": true,
  }

  // Test 1: XML single tool call
  xmlSingle := `<tool_call>
{"name": "search", "arguments": {"query": "golang"}}
</tool_call>`
  calls, rem, ok := DetectToolCalls(xmlSingle, allowedTools)
  if !ok || len(calls) != 1 || calls[0].Function.Name != "search" {
    t.Errorf("Failed to detect XML single tool call: ok=%v, calls=%v, rem=%q", ok, calls, rem)
  }
  if calls[0].Index == nil || *calls[0].Index != 0 {
    t.Errorf("Expected index 0, got %v", calls[0].Index)
  }

  // Test 2: XML multiple tool calls
  xmlMulti := `<tool_call>
{"name": "search", "arguments": {"query": "golang"}}
</tool_call>
<tool_call>
{"name": "weather", "arguments": {"city": "Tokyo"}}
</tool_call>`
  calls2, _, ok2 := DetectToolCalls(xmlMulti, allowedTools)
  if !ok2 || len(calls2) != 2 {
    t.Fatalf("Failed to detect XML multiple tool calls: ok=%v, calls=%v", ok2, calls2)
  }
  if calls2[0].Function.Name != "search" || calls2[1].Function.Name != "weather" {
    t.Errorf("Tool call names mismatch: %v", calls2)
  }
  if calls2[0].Index == nil || *calls2[0].Index != 0 || calls2[1].Index == nil || *calls2[1].Index != 1 {
    t.Errorf("Indices mismatch: %v, %v", calls2[0].Index, calls2[1].Index)
  }

  // Test 3: JSON array inside XML
  xmlArray := `<tool_call>
[
  {"name": "tool_a", "arguments": {"a": 1}},
  {"name": "tool_b", "arguments": {"b": 2}}
]
</tool_call>`
  calls3, _, ok3 := DetectToolCalls(xmlArray, allowedTools)
  if !ok3 || len(calls3) != 2 {
    t.Fatalf("Failed to detect JSON array inside XML tool call: ok=%v, calls=%v", ok3, calls3)
  }
  if calls3[0].Function.Name != "tool_a" || calls3[1].Function.Name != "tool_b" {
    t.Errorf("Tool names mismatch: %v", calls3)
  }

  // Test 4: Direct JSON array without tags
  jsonArray := `[{"name": "direct_tool", "arguments": {"key": "val"}}]`
  calls4, _, ok4 := DetectToolCalls(jsonArray, allowedTools)
  if !ok4 || len(calls4) != 1 || calls4[0].Function.Name != "direct_tool" {
    t.Errorf("Failed to detect direct JSON array: ok=%v, calls=%v", ok4, calls4)
  }
}

func TestWriteCompletionResponseToolCallsNilContent(t *testing.T) {
  rec := httptest.NewRecorder()
  idx := 0
  calls := []ToolCall{
    {
      Index: &idx,
      ID:    "call_test",
      Type:  "function",
      Function: ToolCallFunction{
        Name:      "test_func",
        Arguments: `{}`,
      },
    },
  }

  WriteCompletionResponse(rec, "cmpl-1", 123456, "qwen/qwen3.6-27b", "Some intermediate thought", "", calls, "tool_calls")

  body := rec.Body.String()
  if !strings.Contains(body, `"content":null`) {
    t.Errorf("Expected content:null when tool_calls present, got:\n%s", body)
  }
  if !strings.Contains(body, `"reasoning_content":"Some intermediate thought"`) {
    t.Errorf("Expected preamble text to be preserved in reasoning_content, got:\n%s", body)
  }
  if !strings.Contains(body, `"finish_reason":"tool_calls"`) {
    t.Errorf("Expected finish_reason tool_calls, got:\n%s", body)
  }
}

func TestIsStaticUIMarkdown(t *testing.T) {
  cases := []struct {
    text     string
    prompt   string
    expected bool
  }{
    {"Model: qwen/qwen3.6-27b", "hello", true},
    {"Model: openai/gpt-oss-20b", "hello", true},
    {"Model : deepseek-r1-distill-llama-70b", "hello", true},
    {"- **Model ID:** `qwen/qwen3.6-27b`", "hello", true},
    {"[Powered by Groq](https://groq.com)", "hello", true},
    {"Partial response: generation did not finish.", "hello", true},
    {"[Notice] **Important:** Do not rely on AI without checking", "hello", true},
    {"Important: Do not rely on AI without checking", "hello", true},
    {"3 older messages were left out of this request.", "hello", true},
    {"Here is the solution to your question.", "hello", false},
    {"<think>thinking</think>Result: 42", "hello", false},
    {"<tool_call>{\"name\":\"calc\"}</tool_call>", "hello", false},
  }

  for _, c := range cases {
    got := isStaticUIMarkdown(c.text, c.prompt)
    if got != c.expected {
      t.Errorf("isStaticUIMarkdown(%q) = %v; expected %v", c.text, got, c.expected)
    }
  }
}

func TestFormatPromptCorruptedCaptionFiltering(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen/qwen3.6-27b",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello"},
      {Role: "assistant", Content: "Model: qwen/qwen3.6-27b"},
      {Role: "user", Content: "What is 2+2?"},
    },
  }

  prompt := FormatPrompt(req)
  if strings.Contains(prompt, "Model: qwen/qwen3.6-27b") {
    t.Errorf("Expected corrupted Model: caption to be filtered out of history, got:\n%s", prompt)
  }
  if !strings.Contains(prompt, "What is 2+2?") {
    t.Errorf("Expected user question to remain, got:\n%s", prompt)
  }
}

func TestFormatPromptHistoryPruning(t *testing.T) {
  var messages []ChatMessage
  messages = append(messages, ChatMessage{Role: "user", Content: "Initial user query to remember"})

  // Add 30 turns of long conversation
  for i := 1; i <= 30; i++ {
    messages = append(messages, ChatMessage{Role: "assistant", Content: strings.Repeat("Long assistant response ", 30)})
    messages = append(messages, ChatMessage{Role: "user", Content: strings.Repeat("Long user question ", 30)})
  }
  messages = append(messages, ChatMessage{Role: "user", Content: "Final question"})

  req := ChatCompletionRequest{
    Model:    "qwen/qwen3.6-27b",
    Messages: messages,
  }

  prompt := FormatPrompt(req)

  // Verify initial user query is preserved
  if !strings.Contains(prompt, "Initial user query to remember") {
    t.Errorf("Expected prompt to preserve initial user query")
  }

  // Verify truncation marker is present
  if !strings.Contains(prompt, "[... earlier conversation turns omitted for brevity ...]") {
    t.Errorf("Expected prompt to contain omission notice")
  }

  // Verify prompt size is kept bounded
  if len(prompt) > 20000 {
    t.Errorf("Prompt size too large after pruning: %d chars", len(prompt))
  }
}

func TestParseRetryAfterAndRateLimit(t *testing.T) {
  alert1 := "Rate limit reached. Please retry in a few seconds. Retry after 10 seconds."
  if !isRateLimit(alert1) {
    t.Errorf("Expected alert1 to be identified as rate limit")
  }
  if sec := parseRetryAfter(alert1); sec != 10 {
    t.Errorf("Expected retry seconds 10, got %d", sec)
  }

  alert2 := "429 Too Many Requests: retry in 5"
  if !isRateLimit(alert2) {
    t.Errorf("Expected alert2 to be identified as rate limit")
  }
  if sec := parseRetryAfter(alert2); sec != 5 {
    t.Errorf("Expected retry seconds 5, got %d", sec)
  }

  alert3 := "Your message or image is too large for this model."
  if isRateLimit(alert3) {
    t.Errorf("Expected alert3 NOT to be rate limit")
  }
}

func TestWriteAPIError(t *testing.T) {
  rec := httptest.NewRecorder()
  WriteAPIError(rec, 429, "Rate limit reached. Retry after 10 seconds.", "rate_limit_error", 429)

  if rec.Code != 429 {
    t.Errorf("Expected status 429, got %d", rec.Code)
  }
  body := rec.Body.String()
  if !strings.Contains(body, `"rate_limit_error"`) || !strings.Contains(body, `"code":429`) {
    t.Errorf("Unexpected error response body: %s", body)
  }
}

func TestDetectToolCallsWithUnregisteredTool(t *testing.T) {
  allowed := map[string]bool{"calculator": true}
  content := `Here is an example:
<tool_call>
{"name": "non_existing_tool", "arguments": {"foo": "bar"}}
</tool_call>
Do not execute it.`

  calls, rem, ok := DetectToolCalls(content, allowed)
  if ok || len(calls) > 0 {
    t.Errorf("Expected unregistered tool to not be detected as tool call, got ok=%v, calls=%v", ok, calls)
  }
  if !strings.Contains(rem, "non_existing_tool") {
    t.Errorf("Expected remaining text to keep unregistered tool text intact, got:\n%s", rem)
  }
}

func TestDetectToolCallsWithEmptyAllowedTools(t *testing.T) {
  content := `<tool_call>
{"name": "calculator", "arguments": {"expr": "1+1"}}
</tool_call>`

  calls, rem, ok := DetectToolCalls(content, nil)
  if ok || len(calls) > 0 {
    t.Errorf("Expected tool calls to be disabled when allowedTools is nil, got ok=%v, calls=%v", ok, calls)
  }
  if rem != content {
    t.Errorf("Expected content to remain unchanged when allowedTools is nil")
  }
}

func TestDetectToolCallsWithUnclosedTag(t *testing.T) {
  allowed := map[string]bool{"calculator": true}
  content := "Please explain what <tool_call> means in your prompt. My name is Alice."

  calls, rem, ok := DetectToolCalls(content, allowed)
  if ok || len(calls) > 0 {
    t.Errorf("Expected unclosed tag to not be detected as tool call, got ok=%v, calls=%v", ok, calls)
  }
  if rem != content {
    t.Errorf("Expected content to remain completely unchanged, got:\n%s", rem)
  }
}

func TestDetectToolCallsPartiallyValid(t *testing.T) {
  allowed := map[string]bool{"calculator": true}
  content := `<tool_call>
{"name": "calculator", "arguments": {"expr": "2+2"}}
</tool_call>
And here is a fake tool:
<tool_call>
{"name": "fake_tool", "arguments": {}}
</tool_call>`

  calls, rem, ok := DetectToolCalls(content, allowed)
  if !ok || len(calls) != 1 {
    t.Fatalf("Expected 1 valid tool call, got ok=%v, calls=%v", ok, calls)
  }
  if calls[0].Function.Name != "calculator" {
    t.Errorf("Expected tool name calculator, got %s", calls[0].Function.Name)
  }
  if !strings.Contains(rem, "fake_tool") {
    t.Errorf("Expected remaining content to keep fake_tool block, got:\n%s", rem)
  }
  if strings.Contains(rem, "calculator") {
    t.Errorf("Expected valid calculator block to be removed from remaining content, got:\n%s", rem)
  }
}

func TestStreamToolCallFilterNonExistingTool(t *testing.T) {
  allowed := map[string]bool{"calculator": true}
  filter := NewStreamToolCallFilter(allowed)

  var contentParts []string
  var reasoningParts []string
  var toolCalls []ToolCall

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }
  onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }

  chunks := []string{
    "Here is ",
    "an example:\n",
    "<tool_call>\n",
    "{\"name\": \"unknown_tool\", ",
    "\"arguments\": {}}\n",
    "</tool_call>\n",
    "Done.",
  }

  for _, chunk := range chunks {
    filter.Feed(chunk, onContent, onReasoning, onToolCall)
  }
  filter.Flush(onContent, onReasoning, onToolCall)

  if len(toolCalls) > 0 {
    t.Errorf("Expected zero tool calls emitted for unknown tool, got %d", len(toolCalls))
  }
  if filter.emittedCall {
    t.Errorf("Expected emittedCall to be false")
  }

  fullContent := strings.Join(contentParts, "")
  if !strings.Contains(fullContent, "unknown_tool") {
    t.Errorf("Expected fullContent to contain unknown_tool block, got:\n%s", fullContent)
  }
  if !strings.Contains(fullContent, "Done.") {
    t.Errorf("Expected fullContent to contain Done., got:\n%s", fullContent)
  }
}

func TestStreamToolCallFilterEmptyAllowedTools(t *testing.T) {
  filter := NewStreamToolCallFilter(nil)

  var contentParts []string
  var toolCalls []ToolCall

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onReasoning := func(s string) {}
  onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }

  chunks := []string{"<tool_call>", `{"name":"calc"}`, "</tool_call>"}
  for _, ch := range chunks {
    filter.Feed(ch, onContent, onReasoning, onToolCall)
  }
  filter.Flush(onContent, onReasoning, onToolCall)

  if len(toolCalls) > 0 {
    t.Errorf("Expected zero tool calls when allowedTools is nil")
  }
  fullContent := strings.Join(contentParts, "")
  if fullContent != "<tool_call>{\"name\":\"calc\"}</tool_call>" {
    t.Errorf("Expected direct pass-through, got:\n%s", fullContent)
  }
}

func TestStreamToolCallFilterValidTool(t *testing.T) {
  allowed := map[string]bool{"calculator": true}
  filter := NewStreamToolCallFilter(allowed)

  var contentParts []string
  var reasoningParts []string
  var toolCalls []ToolCall

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }
  onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }

  filter.Feed("Let me calculate that for you.\n", onContent, onReasoning, onToolCall)
  filter.Feed("<tool_call>\n{\"name\": \"calculator\", \"arguments\": {\"expr\": \"5*5\"}}\n</tool_call>", onContent, onReasoning, onToolCall)
  filter.Flush(onContent, onReasoning, onToolCall)

  if len(toolCalls) != 1 {
    t.Fatalf("Expected 1 tool call, got %d", len(toolCalls))
  }
  if toolCalls[0].Function.Name != "calculator" {
    t.Errorf("Expected tool name calculator, got %s", toolCalls[0].Function.Name)
  }
  if !filter.emittedCall {
    t.Errorf("Expected emittedCall to be true")
  }
  fullReasoning := strings.Join(reasoningParts, "")
  if !strings.Contains(fullReasoning, "Let me calculate that") {
    t.Errorf("Expected preamble text in reasoning, got:\n%s", fullReasoning)
  }
}

func TestGetAllowedToolNames(t *testing.T) {
  tools := []Tool{
    {
      Type: "function",
      Function: map[string]interface{}{
        "name": "search",
      },
    },
    {
      Type: "function",
      Function: ToolCallFunction{
        Name: "calculator",
      },
    },
  }

  // Auto choice
  m := GetAllowedToolNames(tools, "auto")
  if !m["search"] || !m["calculator"] || len(m) != 2 {
    t.Errorf("Expected both tools allowed for auto, got %v", m)
  }

  // None choice
  mNone := GetAllowedToolNames(tools, "none")
  if mNone != nil {
    t.Errorf("Expected nil for tool_choice none, got %v", mNone)
  }

  // Specific choice
  mSpecific := GetAllowedToolNames(tools, map[string]interface{}{
    "type": "function",
    "function": map[string]interface{}{
      "name": "calculator",
    },
  })
  if !mSpecific["calculator"] || mSpecific["search"] || len(mSpecific) != 1 {
    t.Errorf("Expected only calculator allowed, got %v", mSpecific)
  }

  // Empty tools
  if GetAllowedToolNames(nil, "auto") != nil {
    t.Errorf("Expected nil for nil tools")
  }
}

func TestFormatPromptAssistantDeduplication(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen/qwen3.6-27b",
    Tools: []Tool{
      {Type: "function", Function: map[string]interface{}{"name": "calc"}},
    },
    Messages: []ChatMessage{
      {Role: "user", Content: "Calculate 2+2"},
      {
        Role:    "assistant",
        Content: `<tool_call>{"name": "calc", "arguments": {"expr": "2+2"}}</tool_call>`,
        ToolCalls: []ToolCall{
          {
            ID:   "call_1",
            Type: "function",
            Function: ToolCallFunction{
              Name:      "calc",
              Arguments: `{"expr": "2+2"}`,
            },
          },
        },
      },
      {Role: "tool", ToolCallID: "call_1", Content: "4"},
    },
  }

  prompt := FormatPrompt(req)

  // Verify there is only one <tool_call> block in the Assistant turn
  assistIdx := strings.Index(prompt, "Assistant:")
  if assistIdx == -1 {
    t.Fatalf("Expected Assistant turn in prompt, got:\n%s", prompt)
  }
  toolExecIdx := strings.Index(prompt, "[Tool Execution Results]")
  assistSection := prompt[assistIdx:toolExecIdx]

  count := strings.Count(assistSection, "<tool_call>")
  if count != 1 {
    t.Errorf("Expected exactly 1 <tool_call> block in Assistant turn, got %d:\n%s", count, assistSection)
  }
}

func TestFormatPromptUserMessageWithToolCall(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen/qwen3.6-27b",
    Messages: []ChatMessage{
      {Role: "user", Content: "How do I format a <tool_call> tag in my script?"},
    },
  }

  prompt := FormatPrompt(req)
  if !strings.Contains(prompt, "How do I format a <tool_call> tag in my script?") {
    t.Errorf("Expected user prompt to preserve <tool_call> text, got:\n%s", prompt)
  }
}