// 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,
    "Qwen/Qwen3.8-Flash-Next": def,
    "qwen3.8-27b":           "Qwen/Qwen3.8-27B",
    "qwen-27b":              "Qwen/Qwen3.8-27B",
    "Qwen/Qwen3.8-27B":      "Qwen/Qwen3.8-27B",
    "qwen3.8-27b-uncensored": "Qwen/Qwen3.8-27B-Uncensored",
    "Qwen/Qwen3.8-27B-Uncensored": "Qwen/Qwen3.8-27B-Uncensored",
    "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([]string{mockServer.URL}, "Qwen/Qwen3.8-Flash-Next", "chat_response", "", "", "", "", true, 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())
  }
}

func TestParseAssistantTextDirectString(t *testing.T) {
  dataJSON := `["Step-by-step reasoning\n</think>\n\nFinal answer here", null]`
  text, ok := parseAssistantText(dataJSON)
  if !ok {
    t.Fatalf("expected successful parse of direct string assistant text")
  }
  if text != "Step-by-step reasoning\n</think>\n\nFinal answer here" {
    t.Fatalf("unexpected text: %q", text)
  }

  reasoning, content := SeparateReasoningAndContent(text)
  if reasoning != "Step-by-step reasoning" {
    t.Fatalf("unexpected reasoning: %q", reasoning)
  }
  if content != "Final answer here" {
    t.Fatalf("unexpected content: %q", content)
  }
}

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

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

      chunk := `event: complete` + "\n" +
        `data: ["Analyzing the problem.\n</think>\n\nThe solution is 42.", null]` + "\n\n"
      w.Write([]byte(chunk))
      flusher.Flush()
      return
    }

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

  svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "respond", "", "", "", "", true, true)

  // Non-streaming test
  rec := httptest.NewRecorder()
  req := ChatCompletionRequest{
    Model: "qwen3.8-27b",
    Messages: []ChatMessage{
      {Role: "user", Content: "What is the answer?"},
    },
    Stream: false,
  }

  err := svc.Chat(rec, nil, req)
  if err != nil {
    t.Fatalf("unexpected error in respond 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 response: %v", err)
  }
  if resp.Choices[0].Message.Content != "The solution is 42." {
    t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
  }
  if resp.Choices[0].Message.ReasoningContent != "Analyzing the problem." {
    t.Fatalf("unexpected reasoning: %v", resp.Choices[0].Message.ReasoningContent)
  }

  // Streaming test
  recStream := httptest.NewRecorder()
  reqStream := ChatCompletionRequest{
    Model: "qwen3.8-27b",
    Messages: []ChatMessage{
      {Role: "user", Content: "What is the answer?"},
    },
    Stream: true,
  }

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

  scanner := bufio.NewScanner(recStream.Body)
  var receivedContent strings.Builder
  var receivedReasoning 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] in stream")
  }
  if receivedReasoning.String() != "Analyzing the problem." {
    t.Fatalf("unexpected streamed reasoning: %q", receivedReasoning.String())
  }
  if receivedContent.String() != "The solution is 42." {
    t.Fatalf("unexpected streamed content: %q", receivedContent.String())
  }
}

func TestQwenServiceDirectOpenAIMock(t *testing.T) {
  mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/v1/chat/completions" && r.Method == http.MethodPost {
      var req ChatCompletionRequest
      if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
        http.Error(w, "bad request", http.StatusBadRequest)
        return
      }

      if !req.Stream {
        w.Header().Set("Content-Type", "application/json")
        json.NewEncoder(w).Encode(ChatCompletionResponse{
          ID:      "chatcmpl-openai-mock",
          Object:  "chat.completion",
          Created: 1234567890,
          Model:   req.Model,
          Choices: []ChatCompletionResponseChoice{
            {
              Index: 0,
              Message: ChatMessage{
                Role:    "assistant",
                Content: "Hello from native OpenAI upstream!",
              },
              FinishReason: "stop",
            },
          },
        })
        return
      }

      // Stream
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)
      chunk := `data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello from "}}]}

data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"native stream!"}}]}

data: [DONE]

`
      w.Write([]byte(chunk))
      if flusher != nil {
        flusher.Flush()
      }
      return
    }

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

  svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "openai", "", "", "", "", true, true)

  // Non-streaming
  rec := httptest.NewRecorder()
  req := ChatCompletionRequest{
    Model: "qwen-openai",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello"},
    },
    Stream: false,
  }
  if err := svc.Chat(rec, nil, req); err != nil {
    t.Fatalf("unexpected error in direct openai 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 response: %v", err)
  }
  if resp.Choices[0].Message.Content != "Hello from native OpenAI upstream!" {
    t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
  }

  // Streaming
  recStream := httptest.NewRecorder()
  reqStream := ChatCompletionRequest{
    Model: "qwen-openai",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello"},
    },
    Stream: true,
  }
  if err := svc.Chat(recStream, nil, reqStream); err != nil {
    t.Fatalf("unexpected error in direct openai streaming: %v", err)
  }
  if !strings.Contains(recStream.Body.String(), "native stream!") {
    t.Fatalf("expected streamed content, got %q", recStream.Body.String())
  }
}

func TestQwenServiceAutoFailover(t *testing.T) {
  // Server 1 simulates ZeroGPU quota exceeded (returns HTTP 429)
  server1Hits := 0
  server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    server1Hits++
    w.WriteHeader(http.StatusTooManyRequests)
    w.Write([]byte(`{"error": "ZeroGPU runs limit reached: quota exceeded"}`))
  }))
  defer server1.Close()

  // Server 2 is healthy and serves completions
  server2Hits := 0
  server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    server2Hits++
    if r.URL.Path == "/v1/chat/completions" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(ChatCompletionResponse{
        ID:      "chatcmpl-failover-mock",
        Object:  "chat.completion",
        Created: 1234567890,
        Choices: []ChatCompletionResponseChoice{
          {
            Index: 0,
            Message: ChatMessage{
              Role:    "assistant",
              Content: "Response from fallback server!",
            },
            FinishReason: "stop",
          },
        },
      })
      return
    }
    http.NotFound(w, r)
  }))
  defer server2.Close()

  endpoints := []string{server1.URL, server2.URL}
  svc := NewQwenService(endpoints, "Qwen/Qwen3.8-27B", "openai", "", "", "", "", true, true)

  rec := httptest.NewRecorder()
  req := ChatCompletionRequest{
    Model: "qwen3.8-27b",
    Messages: []ChatMessage{
      {Role: "user", Content: "Failover test"},
    },
    Stream: false,
  }

  err := svc.Chat(rec, nil, req)
  if err != nil {
    t.Fatalf("expected failover to succeed, got error: %v", err)
  }

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

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

  if resp.Choices[0].Message.Content != "Response from fallback server!" {
    t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
  }

  if server1Hits != 1 {
    t.Fatalf("expected server1 to be attempted once, got %d", server1Hits)
  }
  if server2Hits != 1 {
    t.Fatalf("expected server2 to be called once on failover, got %d", server2Hits)
  }

  // Second request should skip server 1 because it entered cooldown
  rec2 := httptest.NewRecorder()
  err2 := svc.Chat(rec2, nil, req)
  if err2 != nil {
    t.Fatalf("expected second request to succeed directly on server2: %v", err2)
  }
  if server1Hits != 1 {
    t.Fatalf("expected server1 to be skipped during cooldown, but was hit %d times", server1Hits)
  }
  if server2Hits != 2 {
    t.Fatalf("expected server2 to receive the second request, got hits: %d", server2Hits)
  }
}

func TestResolveMaxTokens(t *testing.T) {
  // Default when empty
  if v := ResolveMaxTokens(ChatCompletionRequest{}); v != 4096 {
    t.Fatalf("expected default 4096, got %d", v)
  }

  // Explicit max_tokens
  if v := ResolveMaxTokens(ChatCompletionRequest{MaxTokens: 2048}); v != 2048 {
    t.Fatalf("expected 2048, got %d", v)
  }

  // max_completion_tokens precedence when max_tokens is 0
  if v := ResolveMaxTokens(ChatCompletionRequest{MaxCompletionTokens: 1024}); v != 1024 {
    t.Fatalf("expected 1024, got %d", v)
  }

  // Cap at 32768
  if v := ResolveMaxTokens(ChatCompletionRequest{MaxTokens: 65536}); v != 32768 {
    t.Fatalf("expected cap at 32768, got %d", v)
  }
}