// k3r053n3_test: Test suite for Kimi K3 Gateway
// Created by Luxferre in 2026, released into the public domain

package main

import (
  "bufio"
  "bytes"
  "context"
  "encoding/json"
  "fmt"
  "io"
  "net"
  "net/http"
  "net/http/httptest"
  "strings"
  "testing"
  "time"
)

func TestModelsEndpoint(t *testing.T) {
  svc := NewKimiService("https://test.hf.space", "kimi-k3", "direct:together", "max", "", 8192, 0.7)
  req, w := httptest.NewRequest("GET", "/v1/models", nil), httptest.NewRecorder()
  NewMux(svc).ServeHTTP(w, req)

  if w.Code != http.StatusOK {
    t.Fatalf("expected status 200, got %d", w.Code)
  }
  var res ModelsResponse
  if err := json.NewDecoder(w.Body).Decode(&res); err != nil || len(res.Data) == 0 {
    t.Fatalf("invalid models response: %v", err)
  }
  for _, m := range res.Data {
    if m.ID == "kimi-k3" {
      return
    }
  }
  t.Fatalf("expected kimi-k3 in models list")
}

func TestResolveReasoningEffort(t *testing.T) {
  tests := []struct{ in, def, exp string }{
    {"", "max", "max"}, {"max", "default", "max"}, {"HIGH", "low", "high"},
    {"low", "max", "low"}, {"medium", "max", "default"}, {"default", "max", "default"}, {"unknown", "max", "max"},
  }
  for _, tt := range tests {
    if res := ResolveReasoningEffort(ChatCompletionRequest{ReasoningEffort: tt.in}, tt.def); res != tt.exp {
      t.Errorf("ResolveReasoningEffort(%q, %q) = %q, want %q", tt.in, tt.def, res, tt.exp)
    }
  }
}

func TestResolveBackend(t *testing.T) {
  svc := NewKimiService("https://test.hf.space", "kimi-k3", "direct:together", "max", "", 8192, 0.7)
  tests := []struct{ model, exp string }{
    {"kimi-k3", "direct:together"}, {"kimi-k3:fireworks", "direct:fireworks"},
    {"kimi-k3:together", "direct:together"}, {"kimi-k3:hf-together", "hf:together"},
    {"kimi-k3:hf-fireworks", "hf:fireworks-ai"}, {"kimi-k3:hf-featherless", "hf:featherless-ai"},
    {"kimi-k3:hf-baseten", "hf:baseten"}, {"custom-model", "direct:together"},
  }
  for _, tt := range tests {
    if b := svc.ResolveBackend(tt.model); b != tt.exp {
      t.Errorf("ResolveBackend(%q) = %q, want %q", tt.model, b, tt.exp)
    }
  }
}

func TestParseSOCKS5URL(t *testing.T) {
  tests := []struct {
    in  string
    exp *SOCKS5Config
  }{
    {"", nil},
    {"127.0.0.1:1080", &SOCKS5Config{Address: "127.0.0.1:1080"}},
    {"socks5://127.0.0.1:9050", &SOCKS5Config{Address: "127.0.0.1:9050"}},
    {"socks5h://user:pass@proxy.example.com:1080", &SOCKS5Config{Address: "proxy.example.com:1080", Username: "user", Password: "pass"}},
    {"socks5://myuser@10.0.0.1:1080", &SOCKS5Config{Address: "10.0.0.1:1080", Username: "myuser"}},
    {"127.0.0.1", &SOCKS5Config{Address: "127.0.0.1:1080"}},
  }
  for _, tt := range tests {
    cfg, err := ParseSOCKS5URL(tt.in)
    if err != nil {
      t.Errorf("ParseSOCKS5URL(%q) err: %v", tt.in, err)
    } else if tt.exp == nil && cfg != nil {
      t.Errorf("ParseSOCKS5URL(%q) expected nil, got %+v", tt.in, cfg)
    } else if tt.exp != nil && (cfg.Address != tt.exp.Address || cfg.Username != tt.exp.Username || cfg.Password != tt.exp.Password) {
      t.Errorf("ParseSOCKS5URL(%q) = %+v, want %+v", tt.in, cfg, tt.exp)
    }
  }
}

func startMockSOCKS5Server(t *testing.T, expectedUser, expectedPass string) (string, func()) {
  ln, err := net.Listen("tcp", "127.0.0.1:0")
  if err != nil {
    t.Fatalf("failed to start mock socks5 listener: %v", err)
  }
  stopCh := make(chan struct{})

  go func() {
    for {
      conn, err := ln.Accept()
      if err != nil {
        return
      }
      go func(c net.Conn) {
        defer c.Close()
        hdr := make([]byte, 2)
        if _, err := io.ReadFull(c, hdr); err != nil || hdr[0] != 0x05 {
          return
        }
        methods := make([]byte, int(hdr[1]))
        if _, err := io.ReadFull(c, methods); err != nil {
          return
        }

        if expectedUser != "" {
          c.Write([]byte{0x05, 0x02})
          authHdr := make([]byte, 2)
          if _, err := io.ReadFull(c, authHdr); err != nil {
            return
          }
          userBuf := make([]byte, int(authHdr[1]))
          io.ReadFull(c, userBuf)
          pLenBuf := make([]byte, 1)
          io.ReadFull(c, pLenBuf)
          passBuf := make([]byte, int(pLenBuf[0]))
          io.ReadFull(c, passBuf)

          if string(userBuf) != expectedUser || string(passBuf) != expectedPass {
            c.Write([]byte{0x01, 0x01})
            return
          }
          c.Write([]byte{0x01, 0x00})
        } else {
          c.Write([]byte{0x05, 0x00})
        }

        reqHdr := make([]byte, 4)
        if _, err := io.ReadFull(c, reqHdr); err != nil || reqHdr[0] != 0x05 || reqHdr[1] != 0x01 {
          return
        }
        var targetAddr string
        switch reqHdr[3] {
        case 0x01:
          ipBuf := make([]byte, 4)
          io.ReadFull(c, ipBuf)
          targetAddr = net.IP(ipBuf).String()
        case 0x03:
          lBuf := make([]byte, 1)
          io.ReadFull(c, lBuf)
          dBuf := make([]byte, int(lBuf[0]))
          io.ReadFull(c, dBuf)
          targetAddr = string(dBuf)
        }
        portBuf := make([]byte, 2)
        io.ReadFull(c, portBuf)
        port := (int(portBuf[0]) << 8) | int(portBuf[1])

        remote, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", targetAddr, port), 5*time.Second)
        if err != nil {
          c.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
          return
        }
        defer remote.Close()
        c.Write([]byte{0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 0x1f, 0x90})
        go io.Copy(remote, c)
        io.Copy(c, remote)
      }(conn)
    }
  }()
  return ln.Addr().String(), func() { close(stopCh); ln.Close() }
}

func TestSOCKS5EndToEnd(t *testing.T) {
  mockTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/test" {
      w.Write([]byte("socks5-ok"))
      return
    }
    http.NotFound(w, r)
  }))
  defer mockTarget.Close()

  targetHostPort := strings.TrimPrefix(mockTarget.URL, "http://")
  proxyAddr, cleanupProxy := startMockSOCKS5Server(t, "alice", "secret123")
  defer cleanupProxy()

  proxyURL := fmt.Sprintf("socks5://alice:secret123@%s", proxyAddr)
  ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  defer cancel()

  conn, err := DialSOCKS5(ctx, proxyURL, targetHostPort)
  if err != nil {
    t.Fatalf("DialSOCKS5 failed: %v", err)
  }
  defer conn.Close()

  fmt.Fprintf(conn, "GET /test HTTP/1.1\r\nHost: %s\r\nConnection: close\r\n\r\n", targetHostPort)
  respBody, err := io.ReadAll(conn)
  if err != nil || !strings.Contains(string(respBody), "socks5-ok") {
    t.Fatalf("socks5 direct dial failed: %v, got: %s", err, string(respBody))
  }

  svc := NewKimiService(mockTarget.URL, "kimi-k3", "direct:together", "max", proxyURL, 8192, 0.7)
  req, _ := http.NewRequest("GET", mockTarget.URL+"/test", nil)
  resp, err := svc.client.Do(req)
  if err != nil {
    t.Fatalf("svc.client via SOCKS5 failed: %v", err)
  }
  defer resp.Body.Close()
  body, _ := io.ReadAll(resp.Body)
  if string(body) != "socks5-ok" {
    t.Fatalf("expected 'socks5-ok', got: %s", string(body))
  }
}

func TestParseGradioSnapshot(t *testing.T) {
  // Test 1: Standard reasoning done and answer present with internal </div> tag in reasoning
  sample1 := `[[
    {"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}],"options":null},
    {"role":"assistant","metadata":null,"content":[{"text":"<","type":"text"}],"options":null},
    {"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":">I am thinking about </div> and </tag> tags.</","type":"text"}],"options":null},
    {"role":"assistant","metadata":null,"content":[{"text":">\n\nHello there!","type":"text"}],"options":null}
  ], {"value":{"text":"","files":[]},"__type__":"update"}]`

  content, reasoning, ok := parseGradioSnapshot(sample1)
  if !ok || content != "Hello there!" || reasoning != "I am thinking about </div> and </tag> tags." {
    t.Fatalf("parseGradioSnapshot sample1 failed: content=%q, reasoning=%q", content, reasoning)
  }

  // Test 2: Reasoning still pending
  sample2 := `[[
    {"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}],"options":null},
    {"role":"assistant","metadata":null,"content":[{"text":"<","type":"text"}],"options":null},
    {"role":"assistant","metadata":{"title":"Reasoning","status":"pending"},"content":[{"text":">Still reasoning","type":"text"}],"options":null}
  ], {"value":{"text":"","files":[]},"__type__":"update"}]`

  content2, reasoning2, ok2 := parseGradioSnapshot(sample2)
  if !ok2 || content2 != "" || reasoning2 != "Still reasoning" {
    t.Fatalf("parseGradioSnapshot sample2 failed: content=%q, reasoning=%q", content2, reasoning2)
  }

  // Test 3: Interleaved messages
  sample3 := `[[
    {"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}]},
    {"role":"assistant","metadata":null,"content":[{"text":"<"}]},
    {"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":">Part 1"}]},
    {"role":"assistant","metadata":null,"content":[{"text":"intro text: "}]},
    {"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":". Part 2</"}]},
    {"role":"assistant","metadata":null,"content":[{"text":">\n\n<tool_call>{\"name\":\"fn\"}</tool_call>"}]}
  ]]`
  content3, reasoning3, ok3 := parseGradioSnapshot(sample3)
  if !ok3 || !strings.Contains(content3, "<tool_call>") || !strings.Contains(reasoning3, "Part 1. Part 2") {
    t.Fatalf("parseGradioSnapshot sample3 failed: content=%q, reasoning=%q", content3, reasoning3)
  }
}

func TestDetectToolCalls(t *testing.T) {
  // 1. Standard XML tag with conversational text
  text1 := "Here is the weather:\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Tokyo\"}}\n</tool_call>\nDone."
  tcs1, cleaned1, ok1 := DetectToolCalls(text1)
  if !ok1 || len(tcs1) != 1 || tcs1[0].Function.Name != "get_weather" || !strings.Contains(tcs1[0].Function.Arguments, "Tokyo") {
    t.Fatalf("DetectToolCalls 1 failed: ok=%v, tcs=%+v", ok1, tcs1)
  }
  if strings.Contains(cleaned1, "<tool_call>") || !strings.Contains(cleaned1, "Here is the weather:") {
    t.Fatalf("DetectToolCalls cleaned 1 invalid: %q", cleaned1)
  }

  // 2. tool-call variant
  text2 := "<tool-call>\n{\"name\": \"calc\", \"arguments\": \"{\\\"x\\\": 1}\"}\n</tool-call>"
  tcs2, _, ok2 := DetectToolCalls(text2)
  if !ok2 || len(tcs2) != 1 || tcs2[0].Function.Name != "calc" {
    t.Fatalf("DetectToolCalls 2 failed: ok=%v, tcs=%+v", ok2, tcs2)
  }

  // 3. markdown codeblock format
  text3 := "```tool_call\n{\"name\": \"search\", \"query\": \"golang\"}\n```"
  tcs3, _, ok3 := DetectToolCalls(text3)
  if !ok3 || len(tcs3) != 1 || tcs3[0].Function.Name != "search" {
    t.Fatalf("DetectToolCalls 3 failed: ok=%v, tcs=%+v", ok3, tcs3)
  }

  // 4. unclosed tag at end of response
  text4 := "I will query:\n<tool_call>\n{\"name\": \"fetch_data\", \"id\": 123}"
  tcs4, cleaned4, ok4 := DetectToolCalls(text4)
  if !ok4 || len(tcs4) != 1 || tcs4[0].Function.Name != "fetch_data" || cleaned4 != "I will query:" {
    t.Fatalf("DetectToolCalls 4 unclosed failed: ok=%v, tcs=%+v, cleaned=%q", ok4, tcs4, cleaned4)
  }

  // 5. function call syntax: get_weather(location="Tokyo")
  text5 := "<tool_call>get_weather(location=\"Tokyo\")</tool_call>"
  tcs5, _, ok5 := DetectToolCalls(text5)
  if !ok5 || len(tcs5) != 1 || tcs5[0].Function.Name != "get_weather" || !strings.Contains(tcs5[0].Function.Arguments, "Tokyo") {
    t.Fatalf("DetectToolCalls 5 function syntax failed: ok=%v, tcs=%+v", ok5, tcs5)
  }

  // 6. markdown ```json with tool payload
  text6 := "```json\n{\"name\": \"lookup\", \"arguments\": {\"id\": 42}}\n```"
  tcs6, _, ok6 := DetectToolCalls(text6)
  if !ok6 || len(tcs6) != 1 || tcs6[0].Function.Name != "lookup" {
    t.Fatalf("DetectToolCalls 6 json block failed: ok=%v, tcs=%+v", ok6, tcs6)
  }
}

func TestStreamToolInterceptor(t *testing.T) {
  w := httptest.NewRecorder()
  streamer := NewStreamer(w, nil, "test-id", 123456789, "kimi-k3")
  interceptor := NewStreamToolInterceptor(streamer)

  interceptor.ProcessContentDelta("Checking weather for you...\n<tool_")
  interceptor.ProcessContentDelta("call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"SF\"}}\n</tool_call>\nHave a nice day!")
  interceptor.FlushRemaining()

  if !interceptor.HasToolCalls() {
    t.Fatalf("expected tool call registered in interceptor")
  }
  body := w.Body.String()
  if !strings.Contains(body, "Checking weather for you...") || !strings.Contains(body, "get_weather") || strings.Contains(body, "<tool_call>") {
    t.Fatalf("unexpected stream body: %s", body)
  }
}

func TestEndToEndMockChatCompletion(t *testing.T) {
  mockGradio := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/call/on_submit" && r.Method == "POST" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_12345"})
      return
    }
    if r.URL.Path == "/gradio_api/call/on_submit/evt_12345" && r.Method == "GET" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)
      fmt.Fprintf(w, "event: generating\ndata: [[{\"role\":\"user\",\"metadata\":null,\"content\":[{\"text\":\"Hi\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\"<\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":{\"title\":\"Reasoning\",\"status\":\"done\"},\"content\":[{\"text\":\">Thinking\"}],\"options\":null}]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }
      fmt.Fprintf(w, "event: complete\ndata: [[{\"role\":\"user\",\"metadata\":null,\"content\":[{\"text\":\"Hi\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\"<\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":{\"title\":\"Reasoning\",\"status\":\"done\"},\"content\":[{\"text\":\">Thinking</\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\">\\n\\nHello!\"}],\"options\":null}]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }
      return
    }
    http.NotFound(w, r)
  }))
  defer mockGradio.Close()

  svc := NewKimiService(mockGradio.URL, "kimi-k3", "direct:together", "max", "", 8192, 0.7)
  mux := NewMux(svc)

  // 1. Non-Streaming
  reqB, _ := json.Marshal(ChatCompletionRequest{Model: "kimi-k3", Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, Stream: false})
  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(reqB))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()
  mux.ServeHTTP(w, req)

  var resp ChatCompletionResponse
  if err := json.NewDecoder(w.Body).Decode(&resp); err != nil || len(resp.Choices) == 0 {
    t.Fatalf("non-streaming chat failed: %v", err)
  }
  if resp.Choices[0].Message.ReasoningContent != "Thinking" || resp.Choices[0].Message.GetContentString() != "Hello!" {
    t.Errorf("unexpected choice response: %+v", resp.Choices[0])
  }

  // 2. Streaming
  reqBStream, _ := json.Marshal(ChatCompletionRequest{Model: "kimi-k3", Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, Stream: true})
  reqStream := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(reqBStream))
  reqStream.Header.Set("Content-Type", "application/json")
  wStream := httptest.NewRecorder()
  mux.ServeHTTP(wStream, reqStream)

  scanner := bufio.NewScanner(wStream.Body)
  hasReasoning, hasContent, hasDone := false, false, false
  for scanner.Scan() {
    line := scanner.Text()
    if strings.HasPrefix(line, "data: ") {
      data := strings.TrimPrefix(line, "data: ")
      if data == "[DONE]" {
        hasDone = true
        continue
      }
      var chunk StreamResponse
      if err := json.Unmarshal([]byte(data), &chunk); err == nil && len(chunk.Choices) > 0 {
        if chunk.Choices[0].Delta.ReasoningContent != "" {
          hasReasoning = true
        }
        if chunk.Choices[0].Delta.Content != "" {
          hasContent = true
        }
      }
    }
  }
  if !hasReasoning || !hasContent || !hasDone {
    t.Errorf("streaming test failed: reasoning=%v, content=%v, done=%v", hasReasoning, hasContent, hasDone)
  }
}