// qflash test suite
// Created by Luxferre in 2026, released into the public domain

package main

import (
  "bufio"
  "encoding/json"
  "net/http"
  "net/http/httptest"
  "strings"
  "testing"
)

func TestChatMessageGetContentString(t *testing.T) {
  // String content
  msg1 := ChatMessage{Role: "user", Content: "Hello world"}
  if msg1.GetContentString() != "Hello world" {
    t.Fatalf("expected 'Hello world', got %q", msg1.GetContentString())
  }

  // Multi-part content
  msg2 := ChatMessage{
    Role: "user",
    Content: []interface{}{
      map[string]interface{}{"type": "text", "text": "Part 1 "},
      map[string]interface{}{"type": "text", "text": "Part 2"},
    },
  }
  if msg2.GetContentString() != "Part 1 Part 2" {
    t.Fatalf("expected 'Part 1 Part 2', got %q", msg2.GetContentString())
  }

  // Nil content
  msg3 := ChatMessage{Role: "assistant", Content: nil}
  if msg3.GetContentString() != "" {
    t.Fatalf("expected empty string, got %q", msg3.GetContentString())
  }
}

func TestSOCKS5Parsing(t *testing.T) {
  cfg, err := ParseSOCKS5URL("socks5://user:pass@127.0.0.1:9050")
  if err != nil {
    t.Fatalf("unexpected error: %v", err)
  }
  if cfg.Address != "127.0.0.1:9050" || cfg.Username != "user" || cfg.Password != "pass" {
    t.Fatalf("mismatched parsed socks5 config: %+v", cfg)
  }

  cfg2, err := ParseSOCKS5URL("socks5h://proxy.internal:1080")
  if err != nil {
    t.Fatalf("unexpected error: %v", err)
  }
  if cfg2.Address != "proxy.internal:1080" || cfg2.Username != "" {
    t.Fatalf("mismatched parsed socks5 config: %+v", cfg2)
  }

  cfg3, err := ParseSOCKS5URL("10.0.0.5")
  if err != nil {
    t.Fatalf("unexpected error: %v", err)
  }
  if cfg3.Address != "10.0.0.5:1080" {
    t.Fatalf("expected default port 1080, got %s", cfg3.Address)
  }
}

func TestEffectiveModelID(t *testing.T) {
  def := "Qwen/Qwen3.8-Flash-Next"
  cases := map[string]string{
    "":                      def,
    "qwen3.8-flash-next":    def,
    "qwen-flash-next":       def,
    "qwen-flash":            def,
    "qwen3.8-flash":         def,
    "Qwen/Qwen3.8-Flash-Next": def,
    "custom-org/my-model":   "custom-org/my-model",
  }

  for in, exp := range cases {
    res := EffectiveModelID(in, def)
    if res != exp {
      t.Errorf("EffectiveModelID(%q) = %q; expected %q", in, res, exp)
    }
  }
}

func TestSeparateReasoningAndContentThinkTags(t *testing.T) {
  raw := "<think>\nAnalyzing prompt step by step.\n</think>\n\nHere is the answer."
  reasoning, content := SeparateReasoningAndContent(raw)
  if reasoning != "Analyzing prompt step by step." {
    t.Fatalf("unexpected reasoning: %q", reasoning)
  }
  if content != "Here is the answer." {
    t.Fatalf("unexpected content: %q", content)
  }
}

func TestSeparateReasoningAndContentGradioFormat(t *testing.T) {
  raw := "> 💭 **Thinking Process (QSA Micro-block Reasoning):**\n>\n> Thinking Process:\n>\n> 1. Step one\n> 2. Step two\n\n---\n\n### Answer Header\n\nDetailed answer here."
  reasoning, content := SeparateReasoningAndContent(raw)
  if !strings.Contains(reasoning, "1. Step one") || !strings.Contains(reasoning, "2. Step two") {
    t.Fatalf("expected reasoning to contain steps, got: %q", reasoning)
  }
  if strings.Contains(reasoning, ">") {
    t.Fatalf("expected blockquote markers to be stripped, got: %q", reasoning)
  }
  if content != "### Answer Header\n\nDetailed answer here." {
    t.Fatalf("unexpected content: %q", content)
  }
}

func TestSeparateReasoningAndContentStreamingDivider(t *testing.T) {
  raw := "> 💭 **Thinking Process (QSA Micro-block Reasoning):**\n>\n> Thinking Process:\n>\n> 1. Formulating response...\n\n---\n*Generating response...*"
  reasoning, content := SeparateReasoningAndContent(raw)
  if !strings.Contains(reasoning, "1. Formulating response...") {
    t.Fatalf("expected reasoning, got: %q", reasoning)
  }
  if content != "" {
    t.Fatalf("expected empty content during thought phase, got: %q", content)
  }
}

func TestToolCallParsingAndDetection(t *testing.T) {
  rawJSON := `{"name": "get_weather", "arguments": {"city": "Tokyo"}}`
  tc, ok := parseSingleToolCall(rawJSON)
  if !ok {
    t.Fatalf("expected successful single tool call parse")
  }
  if tc.Function.Name != "get_weather" {
    t.Fatalf("expected 'get_weather', got %q", tc.Function.Name)
  }

  rawXML := `<tool_call>
{"name": "fetch_data", "arguments": "{\"id\": 42}"}
</tool_call>`
  calls, rem, hasCalls := DetectToolCalls(rawXML)
  if !hasCalls || len(calls) != 1 {
    t.Fatalf("expected 1 detected tool call, got %d", len(calls))
  }
  if calls[0].Function.Name != "fetch_data" {
    t.Fatalf("expected 'fetch_data', got %q", calls[0].Function.Name)
  }
  if rem != "" {
    t.Fatalf("expected empty remaining content, got %q", rem)
  }
}

func TestStreamToolCallFilterNoLeak(t *testing.T) {
  filter := &StreamToolCallFilter{}
  var streamedContent strings.Builder
  var emittedCalls []ToolCall

  onContent := func(s string) {
    streamedContent.WriteString(s)
  }
  onTool := func(tc ToolCall) {
    emittedCalls = append(emittedCalls, tc)
  }

  // Stream in small split chunks that split the <tool_call> tag
  chunks := []string{
    "Here is the data: ",
    "<tool",
    "_call>\n",
    `{"name": "query_db", "arguments": {"sql": "SELECT 1"}}`,
    "\n</tool",
    "_call>",
  }

  for _, c := range chunks {
    filter.Feed(c, onContent, onTool)
  }
  filter.Flush(onContent, onTool)

  if strings.Contains(streamedContent.String(), "<tool_call>") || strings.Contains(streamedContent.String(), "</tool_call>") {
    t.Fatalf("tool call tags leaked into content: %q", streamedContent.String())
  }
  if streamedContent.String() != "Here is the data: " {
    t.Fatalf("unexpected content: %q", streamedContent.String())
  }
  if len(emittedCalls) != 1 {
    t.Fatalf("expected 1 emitted tool call, got %d", len(emittedCalls))
  }
  if emittedCalls[0].Function.Name != "query_db" {
    t.Fatalf("expected 'query_db', got %q", emittedCalls[0].Function.Name)
  }
}

func TestParseAssistantText(t *testing.T) {
  dataJSON := `[[
    {"role": "user", "metadata": null, "content": [{"text": "hi", "type": "text"}], "options": null},
    {"role": "assistant", "metadata": null, "content": [{"text": "Hello, human!", "type": "text"}], "options": null}
  ]]`

  text, ok := parseAssistantText(dataJSON)
  if !ok {
    t.Fatalf("expected successful parse of assistant text")
  }
  if text != "Hello, human!" {
    t.Fatalf("expected 'Hello, human!', got %q", text)
  }
}

func TestQwenServiceChatMock(t *testing.T) {
  // Mock upstream Gradio space server
  mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/call/chat_response" && r.Method == http.MethodPost {
      w.Header().Set("Content-Type", "application/json")
      w.Write([]byte(`{"event_id": "test_event_123"}`))
      return
    }

    if r.URL.Path == "/gradio_api/call/chat_response/test_event_123" && r.Method == http.MethodGet {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)

      // Step 1: Thinking progress
      chunk1 := `event: generating` + "\n" +
        `data: [[{"role": "user", "content": [{"text": "Hello", "type": "text"}]}, {"role": "assistant", "content": [{"text": "> 💭 **Thinking Process:**\n>\n> Thinking Process:\n>\n> 1. Step 1\n\n---\n*Generating response...*", "type": "text"}]}]]` + "\n\n"
      w.Write([]byte(chunk1))
      flusher.Flush()

      // Step 2: Final completion
      chunk2 := `event: complete` + "\n" +
        `data: [[{"role": "user", "content": [{"text": "Hello", "type": "text"}]}, {"role": "assistant", "content": [{"text": "> 💭 **Thinking Process:**\n>\n> Thinking Process:\n>\n> 1. Step 1\n\n---\n\nGreetings from mock Qwen!", "type": "text"}]}]]` + "\n\n"
      w.Write([]byte(chunk2))
      flusher.Flush()
      return
    }

    http.NotFound(w, r)
  }))
  defer mockServer.Close()

  svc := NewQwenService(mockServer.URL, "Qwen/Qwen3.8-Flash-Next", "", "", "", "", true)

  // 1. Test Non-streaming completion
  rec := httptest.NewRecorder()
  req := ChatCompletionRequest{
    Model: "qwen3.8-flash-next",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello"},
    },
    Stream: false,
  }

  err := svc.Chat(rec, nil, req)
  if err != nil {
    t.Fatalf("unexpected error in Chat non-streaming: %v", err)
  }

  if rec.Code != http.StatusOK {
    t.Fatalf("expected HTTP 200, got %d", rec.Code)
  }

  var resp ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
    t.Fatalf("failed to decode completion response: %v", err)
  }

  if len(resp.Choices) == 0 {
    t.Fatalf("expected choices, got 0")
  }
  if resp.Choices[0].Message.Content != "Greetings from mock Qwen!" {
    t.Fatalf("unexpected message content: %v", resp.Choices[0].Message.Content)
  }
  if !strings.Contains(resp.Choices[0].Message.ReasoningContent, "1. Step 1") {
    t.Fatalf("unexpected reasoning content: %v", resp.Choices[0].Message.ReasoningContent)
  }

  // 2. Test Streaming completion
  recStream := httptest.NewRecorder()
  reqStream := ChatCompletionRequest{
    Model: "qwen-flash",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello"},
    },
    Stream: true,
  }

  errStream := svc.Chat(recStream, nil, reqStream)
  if errStream != nil {
    t.Fatalf("unexpected error in Chat streaming: %v", errStream)
  }

  scanner := bufio.NewScanner(recStream.Body)
  var receivedReasoning strings.Builder
  var receivedContent strings.Builder
  var sawDone bool

  for scanner.Scan() {
    line := scanner.Text()
    if strings.HasPrefix(line, "data: ") {
      payload := strings.TrimPrefix(line, "data: ")
      if payload == "[DONE]" {
        sawDone = true
        continue
      }
      var sResp StreamResponse
      if err := json.Unmarshal([]byte(payload), &sResp); err == nil && len(sResp.Choices) > 0 {
        delta := sResp.Choices[0].Delta
        if delta.ReasoningContent != "" {
          receivedReasoning.WriteString(delta.ReasoningContent)
        }
        if delta.Content != "" {
          receivedContent.WriteString(delta.Content)
        }
      }
    }
  }

  if !sawDone {
    t.Fatalf("expected [DONE] chunk in stream")
  }
  if !strings.Contains(receivedReasoning.String(), "1. Step 1") {
    t.Fatalf("expected streamed reasoning, got %q", receivedReasoning.String())
  }
  if !strings.Contains(receivedContent.String(), "Greetings from mock Qwen!") {
    t.Fatalf("expected streamed content, got %q", receivedContent.String())
  }
}