// k3r053n3: Standalone OpenAI-compatible gateway for Kimi K3 (cw-105-kimi-k3-gguf-demo)
// Created by Luxferre in 2026, released into the public domain

package main

import (
  "bufio"
  "bytes"
  "context"
  "crypto/rand"
  "encoding/json"
  "flag"
  "fmt"
  "io"
  "net"
  "net/http"
  "os"
  "regexp"
  "strconv"
  "strings"
  "time"
)

var (
  DefaultUserAgent    = "Mozilla/5.0 (X11; Linux x86_64; rv:153.0) Gecko/20100101 Firefox/153.0"
  ConfiguredUserAgent string
  fnCallSyntaxRegex   = regexp.MustCompile(`^([a-zA-Z0-9_\-\.]+)\s*\((.*)\)$`)
  kvPairRegex         = regexp.MustCompile(`([a-zA-Z0-9_]+)\s*=\s*(?:"([^"]*)"|'([^']*)'|([^,\s]+))`)
  toolTagRegexes      = []*regexp.Regexp{
    regexp.MustCompile(`(?s)<tool_call>\s*(.*?)\s*</tool_call>`),
    regexp.MustCompile(`(?s)<tool-call>\s*(.*?)\s*</tool-call>`),
    regexp.MustCompile(`(?s)<function_call>\s*(.*?)\s*</function_call>`),
    regexp.MustCompile(`(?s)<tool>\s*(.*?)\s*</tool>`),
    regexp.MustCompile(`(?s)<action>\s*(.*?)\s*</action>`),
    regexp.MustCompile("(?s)```tool_call\\s*(.*?)\\s*```"),
    regexp.MustCompile("(?s)```tool-call\\s*(.*?)\\s*```"),
    regexp.MustCompile("(?s)```json\\s*(.*?)\\s*```"),
  }
)

// ---------------------------------------------------------------------------
// OpenAI API Data Structures
// ---------------------------------------------------------------------------

type ModelItem struct {
  ID      string `json:"id"`
  Object  string `json:"object"`
  Created int64  `json:"created"`
  OwnedBy string `json:"owned_by"`
}

type ModelsResponse struct {
  Object string      `json:"object"`
  Data   []ModelItem `json:"data"`
}

type ToolCallFunction struct {
  Name      string `json:"name"`
  Arguments string `json:"arguments"`
}

type ToolCall struct {
  Index    *int             `json:"index,omitempty"`
  ID       string           `json:"id,omitempty"`
  Type     string           `json:"type,omitempty"`
  Function ToolCallFunction `json:"function"`
}

type Tool struct {
  Type     string      `json:"type"`
  Function interface{} `json:"function"`
}

type ChatMessage struct {
  Role             string      `json:"role"`
  Content          interface{} `json:"content"`
  ReasoningContent string      `json:"reasoning_content,omitempty"`
  ToolCalls        []ToolCall  `json:"tool_calls,omitempty"`
  ToolCallID       string      `json:"tool_call_id,omitempty"`
  Name             string      `json:"name,omitempty"`
}

func (m *ChatMessage) GetContentString() string {
  switch v := m.Content.(type) {
  case string:
    return v
  case []interface{}:
    var sb strings.Builder
    for _, p := range v {
      if s, ok := p.(string); ok {
        sb.WriteString(s)
      } else if tm, ok := p.(map[string]interface{}); ok {
        if t, ok := tm["text"].(string); ok {
          sb.WriteString(t)
        }
      }
    }
    return sb.String()
  default:
    if m.Content == nil {
      return ""
    }
    b, _ := json.Marshal(m.Content)
    return string(b)
  }
}

type ChatCompletionRequest struct {
  Model               string        `json:"model"`
  Messages            []ChatMessage `json:"messages"`
  Tools               []Tool        `json:"tools,omitempty"`
  Functions           []interface{} `json:"functions,omitempty"`
  ToolChoice          interface{}   `json:"tool_choice,omitempty"`
  FunctionCall        interface{}   `json:"function_call,omitempty"`
  Stream              bool          `json:"stream"`
  MaxTokens           int           `json:"max_tokens"`
  MaxCompletionTokens int           `json:"max_completion_tokens"`
  Temperature         *float64      `json:"temperature,omitempty"`
  ReasoningEffort     string        `json:"reasoning_effort,omitempty"`
}

type ChatCompletionResponseChoice struct {
  Index        int         `json:"index"`
  Message      ChatMessage `json:"message"`
  FinishReason string      `json:"finish_reason"`
}

type Usage struct {
  PromptTokens     int `json:"prompt_tokens"`
  CompletionTokens int `json:"completion_tokens"`
  TotalTokens      int `json:"total_tokens"`
}

type ChatCompletionResponse struct {
  ID      string                         `json:"id"`
  Object  string                         `json:"object"`
  Created int64                          `json:"created"`
  Model   string                         `json:"model"`
  Choices []ChatCompletionResponseChoice `json:"choices"`
  Usage   Usage                          `json:"usage"`
}

type StreamDelta struct {
  Role             string     `json:"role,omitempty"`
  Content          string     `json:"content,omitempty"`
  ReasoningContent string     `json:"reasoning_content,omitempty"`
  ToolCalls        []ToolCall `json:"tool_calls,omitempty"`
}

type StreamChoice struct {
  Index        int         `json:"index"`
  Delta        StreamDelta `json:"delta"`
  FinishReason *string     `json:"finish_reason,omitempty"`
}

type StreamResponse struct {
  ID      string         `json:"id"`
  Object  string         `json:"object"`
  Created int64          `json:"created"`
  Model   string         `json:"model"`
  Choices []StreamChoice `json:"choices"`
}

type GradioJoinResponse struct {
  EventID string `json:"event_id"`
}

// ---------------------------------------------------------------------------
// Zero-Dependency Pure Go SOCKS5 Proxy Client (RFC 1928 / RFC 1929)
// ---------------------------------------------------------------------------

type SOCKS5Config struct {
  Address, Username, Password string
}

func ParseSOCKS5URL(proxyURL string) (*SOCKS5Config, error) {
  u := strings.TrimSpace(proxyURL)
  if u == "" {
    return nil, nil
  }
  u = strings.TrimPrefix(strings.TrimPrefix(u, "socks5h://"), "socks5://")
  cfg := &SOCKS5Config{}
  if at := strings.LastIndex(u, "@"); at != -1 {
    userPass := u[:at]
    cfg.Address = u[at+1:]
    if col := strings.Index(userPass, ":"); col != -1 {
      cfg.Username, cfg.Password = userPass[:col], userPass[col+1:]
    } else {
      cfg.Username = userPass
    }
  } else {
    cfg.Address = u
  }
  if !strings.Contains(cfg.Address, ":") {
    cfg.Address += ":1080"
  }
  return cfg, nil
}

func DialSOCKS5(ctx context.Context, proxyURL, targetAddr string) (net.Conn, error) {
  cfg, err := ParseSOCKS5URL(proxyURL)
  if err != nil || cfg == nil {
    return (&net.Dialer{}).DialContext(ctx, "tcp", targetAddr)
  }
  conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", cfg.Address)
  if err != nil {
    return nil, fmt.Errorf("socks5 dial failed: %w", err)
  }
  if d, ok := ctx.Deadline(); ok {
    conn.SetDeadline(d)
  } else {
    conn.SetDeadline(time.Now().Add(30 * time.Second))
  }
  defer conn.SetDeadline(time.Time{})

  greeting := []byte{0x05, 0x01, 0x00}
  if cfg.Username != "" {
    greeting = []byte{0x05, 0x02, 0x00, 0x02}
  }
  if _, err := conn.Write(greeting); err != nil {
    conn.Close()
    return nil, err
  }
  resp := make([]byte, 2)
  if _, err := io.ReadFull(conn, resp); err != nil || resp[0] != 0x05 {
    conn.Close()
    return nil, fmt.Errorf("socks5 greeting failed")
  }

  if resp[1] == 0x02 {
    req := append(append([]byte{0x01, byte(len(cfg.Username))}, cfg.Username...), byte(len(cfg.Password)))
    req = append(req, cfg.Password...)
    if _, err := conn.Write(req); err != nil {
      conn.Close()
      return nil, err
    }
    if _, err := io.ReadFull(conn, resp); err != nil || resp[1] != 0x00 {
      conn.Close()
      return nil, fmt.Errorf("socks5 auth failed")
    }
  } else if resp[1] != 0x00 {
    conn.Close()
    return nil, fmt.Errorf("socks5 auth rejected: 0x%02x", resp[1])
  }

  host, portStr, err := net.SplitHostPort(targetAddr)
  if err != nil {
    conn.Close()
    return nil, err
  }
  port, _ := strconv.Atoi(portStr)

  req := []byte{0x05, 0x01, 0x00}
  ip := net.ParseIP(host)
  if ip4 := ip.To4(); ip4 != nil {
    req = append(append(req, 0x01), ip4...)
  } else if ip6 := ip.To16(); ip6 != nil {
    req = append(append(req, 0x04), ip6...)
  } else {
    req = append(append(req, 0x03, byte(len(host))), host...)
  }
  req = append(req, byte(port>>8), byte(port&0xFF))

  if _, err := conn.Write(req); err != nil {
    conn.Close()
    return nil, err
  }

  respHdr := make([]byte, 4)
  if _, err := io.ReadFull(conn, respHdr); err != nil || respHdr[1] != 0x00 {
    conn.Close()
    return nil, fmt.Errorf("socks5 connect failed: 0x%02x", respHdr[1])
  }

  switch respHdr[3] {
  case 0x01:
    io.ReadFull(conn, make([]byte, 6))
  case 0x03:
    lb := make([]byte, 1)
    io.ReadFull(conn, lb)
    io.ReadFull(conn, make([]byte, int(lb[0])+2))
  case 0x04:
    io.ReadFull(conn, make([]byte, 18))
  }
  return conn, nil
}

// ---------------------------------------------------------------------------
// Helpers & Tool Calling
// ---------------------------------------------------------------------------

func GenerateUUID() string {
  var b [16]byte
  rand.Read(b[:])
  b[6], b[8] = (b[6]&0x0f)|0x40, (b[8]&0x3f)|0x80
  return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}

func DoWithFibonacciRetry(client *http.Client, makeReq func() (*http.Request, error), maxRetries int) (*http.Response, error) {
  var lastErr error
  a, b := 1, 1
  for attempt := 1; attempt <= maxRetries; attempt++ {
    req, err := makeReq()
    if err != nil {
      return nil, err
    }
    resp, err := client.Do(req)
    if err == nil && resp.StatusCode == http.StatusOK {
      return resp, nil
    }
    if resp != nil {
      body, _ := io.ReadAll(resp.Body)
      resp.Body.Close()
      lastErr = fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
    } else {
      lastErr = err
    }
    if attempt < maxRetries {
      time.Sleep(time.Duration(a) * time.Second)
      a, b = b, a+b
    }
  }
  return nil, fmt.Errorf("request failed after %d retries: %v", maxRetries, lastErr)
}

func EnableCORS(w http.ResponseWriter) {
  w.Header().Set("Access-Control-Allow-Origin", "*")
  w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
  w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, api-key, X-User-Agent")
}

func ResolveMaxTokens(req ChatCompletionRequest, def int) int {
  mt := req.MaxTokens
  if mt == 0 && req.MaxCompletionTokens > 0 {
    mt = req.MaxCompletionTokens
  }
  if mt <= 0 {
    mt = def
  }
  if mt < 256 {
    return 256
  }
  if mt > 8192 {
    return 8192
  }
  return mt
}

func ResolveTemperature(req ChatCompletionRequest, def float64) float64 {
  if req.Temperature != nil {
    t := *req.Temperature
    if t < 0.0 {
      return 0.0
    }
    if t > 1.5 {
      return 1.5
    }
    return t
  }
  return def
}

func ResolveReasoningEffort(req ChatCompletionRequest, def string) string {
  e := strings.ToLower(strings.TrimSpace(req.ReasoningEffort))
  if e == "" {
    e = strings.ToLower(strings.TrimSpace(def))
  }
  switch e {
  case "max", "high", "low":
    return e
  case "default", "medium", "none", "off":
    return "default"
  default:
    return "max"
  }
}

func EffectiveUserAgent(r *http.Request) string {
  if r != nil {
    if c := r.Header.Get("X-User-Agent"); c != "" {
      return c
    }
  }
  if ConfiguredUserAgent != "" {
    return ConfiguredUserAgent
  }
  return DefaultUserAgent
}

func formatArgumentsString(args interface{}) string {
  if args == nil {
    return "{}"
  }
  if str, ok := args.(string); ok {
    return str
  }
  b, err := json.Marshal(args)
  if err != nil {
    return "{}"
  }
  return string(b)
}

func parseFnArgsString(raw string) string {
  raw = strings.TrimSpace(raw)
  if raw == "" {
    return "{}"
  }
  var m map[string]interface{}
  if err := json.Unmarshal([]byte(raw), &m); err == nil {
    b, _ := json.Marshal(m)
    return string(b)
  }
  kvMap := make(map[string]interface{})
  for _, match := range kvPairRegex.FindAllStringSubmatch(raw, -1) {
    k := match[1]
    v := match[2]
    if v == "" {
      v = match[3]
    }
    if v == "" {
      v = match[4]
    }
    kvMap[k] = v
  }
  if len(kvMap) > 0 {
    b, _ := json.Marshal(kvMap)
    return string(b)
  }
  b, _ := json.Marshal(map[string]string{"input": raw})
  return string(b)
}

func BuildToolsSystemPrompt(tools []Tool, toolChoice interface{}) string {
  if len(tools) == 0 {
    return ""
  }
  b, _ := json.MarshalIndent(tools, "", "  ")

  choiceInstruction := ""
  if tcStr, ok := toolChoice.(string); ok {
    if tcStr == "required" {
      choiceInstruction = "\nIMPORTANT: You MUST invoke at least one tool to satisfy the request."
    }
  } else if tcMap, ok := toolChoice.(map[string]interface{}); ok {
    if fn, ok := tcMap["function"].(map[string]interface{}); ok {
      if fnName, ok := fn["name"].(string); ok && fnName != "" {
        choiceInstruction = fmt.Sprintf("\nIMPORTANT: You MUST call the %q tool.", fnName)
      }
    }
  }

  return fmt.Sprintf("# Tools Available\nYou have access to the following tools:\n%s\n\n# Tool Calling Instructions\n1. When a task requires gathering information, inspecting files, running commands, or executing actions, you MUST emit the <tool_call>...</tool_call> block in your response.\n2. DO NOT output conversational promises or filler statements (e.g. \"I will analyze the project\", \"Let me check the files\") without emitting the tool call.\n3. Wrap tool calls inside <tool_call>...</tool_call> XML tags:\n<tool_call>\n{\"name\": \"tool_name\", \"arguments\": {\"param\": \"value\"}}\n</tool_call>%s", string(b), choiceInstruction)
}

func extractToolCallFromMap(m map[string]interface{}, idx int) *ToolCall {
  var name string
  var args interface{}

  for _, k := range []string{"name", "tool", "tool_name", "function_name", "action"} {
    if n, ok := m[k].(string); ok && n != "" {
      name = n
      break
    }
  }
  if name == "" {
    if fn, ok := m["function"].(map[string]interface{}); ok {
      for _, k := range []string{"name", "tool", "tool_name", "function_name"} {
        if n, ok := fn[k].(string); ok && n != "" {
          name = n
          args = fn["arguments"]
          if args == nil {
            args = fn["parameters"]
          }
          if args == nil {
            args = fn["args"]
          }
          break
        }
      }
    }
  }

  if name == "" {
    return nil
  }

  if args == nil {
    for _, k := range []string{"arguments", "parameters", "args", "params", "input", "action_input"} {
      if v, exists := m[k]; exists && v != nil {
        args = v
        break
      }
    }
  }

  if args == nil {
    rem := make(map[string]interface{})
    for k, v := range m {
      if k != "name" && k != "function" && k != "tool" && k != "tool_name" && k != "function_name" && k != "action" && k != "type" && k != "id" && k != "index" {
        rem[k] = v
      }
    }
    if len(rem) > 0 {
      args = rem
    } else {
      args = map[string]interface{}{}
    }
  }

  id := fmt.Sprintf("call_%s_%d", GenerateUUID()[:8], idx)
  if existingID, ok := m["id"].(string); ok && existingID != "" {
    id = existingID
  }

  return &ToolCall{
    Index: &idx,
    ID:    id,
    Type:  "function",
    Function: ToolCallFunction{
      Name:      name,
      Arguments: formatArgumentsString(args),
    },
  }
}

func parseRawToolPayload(raw string) []ToolCall {
  raw = strings.TrimSpace(raw)
  if raw == "" {
    return nil
  }

  var arr []map[string]interface{}
  if err := json.Unmarshal([]byte(raw), &arr); err == nil {
    var res []ToolCall
    for _, item := range arr {
      if tc := extractToolCallFromMap(item, len(res)); tc != nil {
        res = append(res, *tc)
      }
    }
    if len(res) > 0 {
      return res
    }
  }

  var single map[string]interface{}
  if err := json.Unmarshal([]byte(raw), &single); err == nil {
    if tcList, ok := single["tool_calls"].([]interface{}); ok {
      var res []ToolCall
      for _, item := range tcList {
        if m, ok := item.(map[string]interface{}); ok {
          if tc := extractToolCallFromMap(m, len(res)); tc != nil {
            res = append(res, *tc)
          }
        }
      }
      if len(res) > 0 {
        return res
      }
    }
    if tc := extractToolCallFromMap(single, 0); tc != nil {
      return []ToolCall{*tc}
    }
  }

  if fnMatch := fnCallSyntaxRegex.FindStringSubmatch(raw); len(fnMatch) >= 3 {
    fnName := fnMatch[1]
    argsRaw := strings.TrimSpace(fnMatch[2])
    idx := 0
    return []ToolCall{{
      Index: &idx,
      ID:    fmt.Sprintf("call_%s_0", GenerateUUID()[:8]),
      Type:  "function",
      Function: ToolCallFunction{
        Name:      fnName,
        Arguments: parseFnArgsString(argsRaw),
      },
    }}
  }

  return nil
}

func DetectToolCalls(text string) ([]ToolCall, string, bool) {
  var toolCalls []ToolCall
  cleaned := text

  for _, re := range toolTagRegexes {
    matches := re.FindAllStringSubmatchIndex(cleaned, -1)
    if len(matches) == 0 {
      continue
    }
    var sb strings.Builder
    lastIdx := 0
    foundInRe := false
    for _, m := range matches {
      raw := strings.TrimSpace(cleaned[m[2]:m[3]])
      parsed := parseRawToolPayload(raw)
      if len(parsed) > 0 {
        foundInRe = true
        sb.WriteString(cleaned[lastIdx:m[0]])
        lastIdx = m[1]
        for _, tc := range parsed {
          idx := len(toolCalls)
          tc.Index = &idx
          toolCalls = append(toolCalls, tc)
        }
      }
    }
    if foundInRe {
      sb.WriteString(cleaned[lastIdx:])
      cleaned = sb.String()
    }
  }

  if len(toolCalls) == 0 {
    for _, tag := range []string{"<tool_call>", "<tool-call>", "<function_call>", "<tool>", "<action>"} {
      if idx := strings.Index(cleaned, tag); idx != -1 {
        raw := strings.TrimSpace(cleaned[idx+len(tag):])
        for _, tc := range parseRawToolPayload(raw) {
          idxTC := len(toolCalls)
          tc.Index = &idxTC
          toolCalls = append(toolCalls, tc)
        }
        if len(toolCalls) > 0 {
          cleaned = strings.TrimSpace(cleaned[:idx])
          break
        }
      }
    }
  }

  return toolCalls, strings.TrimSpace(cleaned), len(toolCalls) > 0
}

// ---------------------------------------------------------------------------
// Streaming & Snapshot Parser
// ---------------------------------------------------------------------------

type Streamer struct {
  w       http.ResponseWriter
  flusher http.Flusher
  id      string
  created int64
  model   string
}

func NewStreamer(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string) *Streamer {
  w.Header().Set("Content-Type", "text/event-stream")
  w.Header().Set("Cache-Control", "no-cache")
  w.Header().Set("Connection", "keep-alive")
  return &Streamer{w: w, flusher: flusher, id: id, created: created, model: model}
}

func (s *Streamer) SendChunk(delta StreamDelta, finish ...string) {
  var finishReason *string
  if len(finish) > 0 && finish[0] != "" {
    finishReason = &finish[0]
  }
  chunk := StreamResponse{
    ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model,
    Choices: []StreamChoice{{Index: 0, Delta: delta, FinishReason: finishReason}},
  }
  b, _ := json.Marshal(chunk)
  fmt.Fprintf(s.w, "data: %s\n\n", b)
  if s.flusher != nil {
    s.flusher.Flush()
  }
}

func (s *Streamer) Done() {
  fmt.Fprintf(s.w, "data: [DONE]\n\n")
  if s.flusher != nil {
    s.flusher.Flush()
  }
}

func cleanReasoningArtifacts(s string) string {
  s = strings.TrimPrefix(s, "<>")
  s = strings.TrimPrefix(s, ">")
  s = strings.TrimSuffix(s, "</>")
  s = strings.TrimSuffix(s, "</")
  s = strings.TrimSuffix(s, ">")
  return s
}

func cleanContentArtifacts(s string) string {
  if s == "…" || s == "<" || s == ">" || s == "<>" || s == "</" || s == "</>" {
    return ""
  }
  s = strings.TrimPrefix(s, ">\n\n")
  s = strings.TrimPrefix(s, ">\n")
  s = strings.TrimPrefix(s, ">")
  s = strings.TrimPrefix(s, "<>")
  s = strings.TrimSuffix(s, "</>")
  s = strings.TrimSuffix(s, "</")
  return s
}

func parseGradioSnapshot(dataJSON string) (string, string, bool) {
  var raw []interface{}
  if err := json.Unmarshal([]byte(dataJSON), &raw); err != nil || len(raw) == 0 {
    return "", "", false
  }
  msgs, ok := raw[0].([]interface{})
  if !ok || len(msgs) == 0 {
    return "", "", false
  }

  lastUserIdx := -1
  for i := len(msgs) - 1; i >= 0; i-- {
    if m, ok := msgs[i].(map[string]interface{}); ok && m["role"] == "user" {
      lastUserIdx = i
      break
    }
  }

  var assistantMsgs []interface{}
  if lastUserIdx != -1 && lastUserIdx+1 < len(msgs) {
    assistantMsgs = msgs[lastUserIdx+1:]
  } else {
    assistantMsgs = msgs
  }

  var reasoningParts []string
  var contentParts []string
  var rawAllText strings.Builder

  for _, item := range assistantMsgs {
    m, ok := item.(map[string]interface{})
    if !ok || m["role"] != "assistant" {
      continue
    }

    var textParts strings.Builder
    if cArr, ok := m["content"].([]interface{}); ok {
      for _, part := range cArr {
        if cMap, ok := part.(map[string]interface{}); ok {
          if rawText, ok := cMap["text"].(string); ok {
            textParts.WriteString(rawText)
          }
        }
      }
    }
    msgText := textParts.String()
    rawAllText.WriteString(msgText)

    if msgText == "…" || msgText == "<" || msgText == ">" || msgText == "<>" || msgText == "</" || msgText == "</>" {
      continue
    }

    isReasoning := false
    if meta, ok := m["metadata"].(map[string]interface{}); ok && meta != nil {
      if t, ok := meta["title"].(string); ok && strings.EqualFold(t, "Reasoning") {
        isReasoning = true
      } else if st, ok := meta["status"].(string); ok && (st == "pending" || st == "done") {
        isReasoning = true
      }
    }

    if !isReasoning {
      if strings.HasPrefix(msgText, "<>") || strings.HasSuffix(msgText, "</") || strings.HasSuffix(msgText, "</>") {
        isReasoning = true
      } else if strings.HasPrefix(msgText, ">") && !strings.HasPrefix(msgText, ">\n") {
        isReasoning = true
      }
    }

    if isReasoning {
      cleaned := cleanReasoningArtifacts(msgText)
      if cleaned != "" {
        reasoningParts = append(reasoningParts, cleaned)
      }
    } else {
      cleaned := cleanContentArtifacts(msgText)
      if cleaned != "" {
        contentParts = append(contentParts, cleaned)
      }
    }
  }

  reasoning := strings.Join(reasoningParts, "")
  content := strings.Join(contentParts, "")

  if reasoning != "" || content != "" {
    return content, reasoning, true
  }

  combined := rawAllText.String()
  if combined == "…" || combined == "<" || combined == ">" || combined == "<>" {
    return "", "", true
  }

  if strings.Contains(combined, "<think>") {
    tIdx := strings.Index(combined, "<think>")
    endIdx := strings.Index(combined, "</think>")
    if endIdx != -1 {
      return strings.TrimLeft(combined[:tIdx]+combined[endIdx+8:], "\n"), combined[tIdx+7 : endIdx], true
    }
    return strings.TrimSpace(combined[:tIdx]), combined[tIdx+7:], true
  }

  if strings.HasPrefix(combined, "<>") {
    inner := combined[2:]
    if idx := strings.Index(inner, "</>"); idx != -1 {
      return strings.TrimLeft(inner[idx+3:], "\n"), inner[:idx], true
    }
    if idx := strings.LastIndex(inner, "</"); idx != -1 {
      return strings.TrimLeft(strings.TrimPrefix(inner[idx+2:], ">"), "\n"), inner[:idx], true
    }
    return "", inner, true
  }

  return cleanContentArtifacts(combined), "", true
}

type StreamToolInterceptor struct {
  streamer       *Streamer
  buffer         string
  inToolTag      bool
  matchedOpen    string
  closingTag     string
  emittedContent string
  toolCallCount  int
}

func NewStreamToolInterceptor(streamer *Streamer) *StreamToolInterceptor {
  return &StreamToolInterceptor{streamer: streamer}
}

func (si *StreamToolInterceptor) HasToolCalls() bool {
  return si.toolCallCount > 0
}

func (si *StreamToolInterceptor) ProcessContentDelta(delta string) {
  if delta == "" {
    return
  }
  si.buffer += delta

  tags := []struct{ open, close string }{
    {"<tool_call>", "</tool_call>"},
    {"<tool-call>", "</tool-call>"},
    {"<function_call>", "</function_call>"},
    {"<tool>", "</tool>"},
    {"<action>", "</action>"},
    {"```tool_call", "```"},
    {"```tool-call", "```"},
    {"```json", "```"},
  }

  for len(si.buffer) > 0 {
    if !si.inToolTag {
      openIdx := -1
      var foundOpen, foundClose string
      for _, t := range tags {
        if idx := strings.Index(si.buffer, t.open); idx != -1 {
          if openIdx == -1 || idx < openIdx {
            openIdx = idx
            foundOpen = t.open
            foundClose = t.close
          }
        }
      }

      if openIdx == -1 {
        maxOverlap := 0
        for _, t := range tags {
          for o := len(t.open) - 1; o > 0; o-- {
            if strings.HasSuffix(si.buffer, t.open[:o]) && o > maxOverlap {
              maxOverlap = o
            }
          }
        }
        if maxOverlap > 0 {
          toEmit := si.buffer[:len(si.buffer)-maxOverlap]
          if toEmit != "" {
            si.streamer.SendChunk(StreamDelta{Content: toEmit})
            si.emittedContent += toEmit
            si.buffer = si.buffer[len(toEmit):]
          }
          return
        }
        si.streamer.SendChunk(StreamDelta{Content: si.buffer})
        si.emittedContent += si.buffer
        si.buffer = ""
        return
      }

      if openIdx > 0 {
        pre := si.buffer[:openIdx]
        si.streamer.SendChunk(StreamDelta{Content: pre})
        si.emittedContent += pre
      }
      si.inToolTag = true
      si.matchedOpen = foundOpen
      si.closingTag = foundClose
      si.buffer = si.buffer[openIdx+len(foundOpen):]
    }

    if si.inToolTag {
      closeIdx := strings.Index(si.buffer, si.closingTag)
      if closeIdx == -1 {
        return
      }
      toolJSON := strings.TrimSpace(si.buffer[:closeIdx])
      si.buffer = si.buffer[closeIdx+len(si.closingTag):]
      si.inToolTag = false
      parsed := parseRawToolPayload(toolJSON)
      if len(parsed) > 0 {
        for _, tc := range parsed {
          idx := si.toolCallCount
          si.toolCallCount++
          tc.Index = &idx
          si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
        }
      } else {
        rawReconstruct := si.matchedOpen + toolJSON + si.closingTag
        si.streamer.SendChunk(StreamDelta{Content: rawReconstruct})
        si.emittedContent += rawReconstruct
      }
    }
  }
}

func (si *StreamToolInterceptor) FlushRemaining() {
  if si.buffer == "" {
    return
  }
  if si.inToolTag {
    raw := strings.TrimSpace(strings.TrimSuffix(si.buffer, si.closingTag))
    tcs := parseRawToolPayload(raw)
    if len(tcs) > 0 {
      for _, tc := range tcs {
        idx := si.toolCallCount
        si.toolCallCount++
        tc.Index = &idx
        si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
      }
      si.buffer = ""
      return
    }
  }
  tcs, cleaned, ok := DetectToolCalls(si.buffer)
  if ok {
    if cleaned != "" {
      si.streamer.SendChunk(StreamDelta{Content: cleaned})
      si.emittedContent += cleaned
    }
    for _, tc := range tcs {
      idx := si.toolCallCount
      si.toolCallCount++
      tc.Index = &idx
      si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
    }
  } else {
    si.streamer.SendChunk(StreamDelta{Content: si.buffer})
    si.emittedContent += si.buffer
  }
  si.buffer = ""
}

// ---------------------------------------------------------------------------
// Kimi K3 Service & Upstream Logic
// ---------------------------------------------------------------------------

type KimiService struct {
  endpoint, defaultModel, defaultBackend, defaultReason, socks5Proxy string
  defaultTokens                                                      int
  defaultTemp                                                        float64
  client                                                             *http.Client
}

func NewKimiService(endpoint, model, backend, reasoning, socks5Proxy string, tokens int, temp float64) *KimiService {
  if model == "" {
    model = "kimi-k3"
  }
  if backend == "" {
    backend = "direct:together"
  }
  if reasoning == "" {
    reasoning = "max"
  }
  if tokens <= 0 {
    tokens = 8192
  }
  if temp <= 0 {
    temp = 0.7
  }
  cleanProxy := strings.TrimSpace(socks5Proxy)
  transport := &http.Transport{
    DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
      if cleanProxy != "" {
        return DialSOCKS5(ctx, cleanProxy, addr)
      }
      return (&net.Dialer{}).DialContext(ctx, network, addr)
    },
    ForceAttemptHTTP2:   true,
    MaxIdleConns:        100,
    IdleConnTimeout:     90 * time.Second,
    TLSHandshakeTimeout: 15 * time.Second,
  }
  return &KimiService{
    endpoint:       strings.TrimRight(endpoint, "/"),
    defaultModel:   model,
    defaultBackend: backend,
    defaultReason:  reasoning,
    socks5Proxy:    cleanProxy,
    defaultTokens:  tokens,
    defaultTemp:    temp,
    client:         &http.Client{Transport: transport, Timeout: 300 * time.Second},
  }
}

func (s *KimiService) ListModels() []ModelItem {
  now := time.Now().Unix()
  return []ModelItem{
    {ID: s.defaultModel, Object: "model", Created: now, OwnedBy: "moonshotai"},
    {ID: "moonshotai/Kimi-K3", Object: "model", Created: now, OwnedBy: "moonshotai"},
    {ID: "kimi-k3:together", Object: "model", Created: now, OwnedBy: "together"},
    {ID: "kimi-k3:fireworks", Object: "model", Created: now, OwnedBy: "fireworks"},
    {ID: "kimi-k3:hf-together", Object: "model", Created: now, OwnedBy: "huggingface"},
    {ID: "kimi-k3:hf-fireworks", Object: "model", Created: now, OwnedBy: "huggingface"},
    {ID: "kimi-k3:hf-featherless", Object: "model", Created: now, OwnedBy: "huggingface"},
    {ID: "kimi-k3:hf-baseten", Object: "model", Created: now, OwnedBy: "huggingface"},
  }
}

func (s *KimiService) ResolveBackend(modelName string) string {
  if strings.Contains(modelName, ":") {
    switch strings.SplitN(modelName, ":", 2)[1] {
    case "together", "direct-together", "direct:together":
      return "direct:together"
    case "fireworks", "direct-fireworks", "direct:fireworks":
      return "direct:fireworks"
    case "hf-together", "hf:together":
      return "hf:together"
    case "hf-fireworks", "hf:fireworks", "hf:fireworks-ai":
      return "hf:fireworks-ai"
    case "hf-featherless", "hf:featherless", "hf:featherless-ai":
      return "hf:featherless-ai"
    case "hf-baseten", "hf:baseten":
      return "hf:baseten"
    }
  }
  return s.defaultBackend
}

func formatToolCallBlock(tc ToolCall) string {
  var argsObj interface{}
  argsRaw := strings.TrimSpace(tc.Function.Arguments)
  if argsRaw == "" {
    argsObj = map[string]interface{}{}
  } else if err := json.Unmarshal([]byte(argsRaw), &argsObj); err != nil {
    argsObj = map[string]string{"input": argsRaw}
  }
  b, _ := json.Marshal(map[string]interface{}{
    "name":      tc.Function.Name,
    "arguments": argsObj,
  })
  return fmt.Sprintf("<tool_call>\n%s\n</tool_call>", string(b))
}

func formatMessageContent(m ChatMessage) (string, string) {
  role := m.Role
  cStr := m.GetContentString()
  if role == "tool" || role == "function" {
    tName := m.Name
    if tName == "" {
      tName = m.ToolCallID
    }
    if tName != "" {
      cStr = fmt.Sprintf("[Tool Result for %s]: %s", tName, cStr)
    } else {
      cStr = fmt.Sprintf("[Tool Result]: %s", cStr)
    }
    role = "user"
  } else if role == "assistant" && len(m.ToolCalls) > 0 {
    var tcParts []string
    for _, tc := range m.ToolCalls {
      tcParts = append(tcParts, formatToolCallBlock(tc))
    }
    if cStr != "" {
      cStr += "\n\n" + strings.Join(tcParts, "\n")
    } else {
      cStr = strings.Join(tcParts, "\n")
    }
  }
  return role, cStr
}

func (s *KimiService) Chat(w http.ResponseWriter, r *http.Request, req ChatCompletionRequest) error {
  modelName := req.Model
  if modelName == "" {
    modelName = s.defaultModel
  }

  if len(req.Tools) == 0 && len(req.Functions) > 0 {
    for _, fn := range req.Functions {
      req.Tools = append(req.Tools, Tool{Type: "function", Function: fn})
    }
    if req.ToolChoice == nil && req.FunctionCall != nil {
      req.ToolChoice = req.FunctionCall
    }
  }

  var nonSys []ChatMessage
  var sysParts []string
  for _, m := range req.Messages {
    if m.Role == "system" {
      if c := m.GetContentString(); c != "" {
        sysParts = append(sysParts, c)
      }
    } else {
      nonSys = append(nonSys, m)
    }
  }
  userSysPrompt := strings.Join(sysParts, "\n\n")

  toolsPrompt := ""
  if len(req.Tools) > 0 {
    toolsPrompt = BuildToolsSystemPrompt(req.Tools, req.ToolChoice)
  }

  var history []map[string]interface{}
  var promptText string

  if len(nonSys) == 0 {
    promptText = "Hello"
    if toolsPrompt != "" {
      promptText = toolsPrompt + "\n\n" + promptText
    }
    if userSysPrompt != "" {
      promptText = userSysPrompt + "\n\n" + promptText
    }
  } else {
    firstTrailingTool := -1
    for i := len(nonSys) - 1; i >= 0; i-- {
      if nonSys[i].Role == "tool" || nonSys[i].Role == "function" {
        firstTrailingTool = i
      } else {
        break
      }
    }

    var historyMsgs []ChatMessage
    var currentTurnMsgs []ChatMessage

    if firstTrailingTool != -1 {
      historyMsgs = nonSys[:firstTrailingTool]
      currentTurnMsgs = nonSys[firstTrailingTool:]
    } else {
      historyMsgs = nonSys[:len(nonSys)-1]
      currentTurnMsgs = nonSys[len(nonSys)-1:]
    }

    for i, m := range historyMsgs {
      role, cStr := formatMessageContent(m)
      if i == 0 && role == "user" {
        if userSysPrompt != "" {
          cStr = userSysPrompt + "\n\n" + cStr
        }
        if toolsPrompt != "" {
          cStr = toolsPrompt + "\n\n" + cStr
        }
      }
      history = append(history, map[string]interface{}{
        "role": role, "metadata": nil,
        "content": []map[string]interface{}{{"text": cStr, "type": "text"}},
        "options": nil,
      })
    }

    if firstTrailingTool != -1 {
      var toolResParts []string
      for _, m := range currentTurnMsgs {
        _, c := formatMessageContent(m)
        toolResParts = append(toolResParts, c)
      }
      promptText = strings.Join(toolResParts, "\n\n") + "\n\nBased on the tool results above, continue your task. If you need to invoke another tool, respond with <tool_call>{\"name\": \"tool_name\", \"arguments\": {...}}</tool_call>. Otherwise, provide your final response."
    } else {
      _, promptText = formatMessageContent(currentTurnMsgs[0])
    }

    if toolsPrompt != "" {
      promptText = toolsPrompt + "\n\n" + promptText
    }
    if len(history) == 0 && userSysPrompt != "" {
      promptText = userSysPrompt + "\n\n" + promptText
    }
  }

  gradioPayload, _ := json.Marshal(map[string]interface{}{
    "data": []interface{}{
      map[string]interface{}{"text": promptText, "files": []interface{}{}},
      history,
      s.ResolveBackend(modelName),
      ResolveMaxTokens(req, s.defaultTokens),
      ResolveTemperature(req, s.defaultTemp),
      ResolveReasoningEffort(req, s.defaultReason),
    },
  })

  effUA := EffectiveUserAgent(r)
  resp, err := DoWithFibonacciRetry(s.client, func() (*http.Request, error) {
    req, err := http.NewRequest("POST", s.endpoint+"/gradio_api/call/on_submit", bytes.NewBuffer(gradioPayload))
    if err == nil {
      req.Header.Set("Content-Type", "application/json")
      req.Header.Set("User-Agent", effUA)
    }
    return req, err
  }, 5)
  if err != nil {
    return fmt.Errorf("upstream error: %w", err)
  }
  defer resp.Body.Close()

  var joinRes GradioJoinResponse
  if err := json.NewDecoder(resp.Body).Decode(&joinRes); err != nil || joinRes.EventID == "" {
    return fmt.Errorf("failed to parse Gradio event ID")
  }

  streamResp, err := DoWithFibonacciRetry(s.client, func() (*http.Request, error) {
    req, err := http.NewRequest("GET", fmt.Sprintf("%s/gradio_api/call/on_submit/%s", s.endpoint, joinRes.EventID), nil)
    if err == nil {
      req.Header.Set("Accept", "text/event-stream")
      req.Header.Set("User-Agent", effUA)
    }
    return req, err
  }, 5)
  if err != nil {
    return fmt.Errorf("upstream stream error: %w", err)
  }
  defer streamResp.Body.Close()

  compID := "chatcmpl-" + GenerateUUID()
  created := time.Now().Unix()

  if !req.Stream {
    reader := bufio.NewReader(streamResp.Body)
    var finalContent, finalReasoning string
    for {
      line, err := reader.ReadString('\n')
      if err != nil {
        break
      }
      line = strings.TrimSpace(line)
      if strings.HasPrefix(line, "data: ") {
        if c, r, ok := parseGradioSnapshot(strings.TrimPrefix(line, "data: ")); ok {
          if c != "" {
            finalContent = c
          }
          if r != "" {
            finalReasoning = r
          }
        }
      }
    }

    toolCalls, cleanedContent, hasTools := DetectToolCalls(finalContent)
    if !hasTools && len(req.Tools) > 0 {
      if rCalls, _, rHas := DetectToolCalls(finalReasoning); rHas {
        toolCalls = rCalls
        hasTools = true
      }
    }
    if !hasTools && len(req.Tools) > 0 {
      if directTCs := parseRawToolPayload(finalContent); len(directTCs) > 0 {
        toolCalls = directTCs
        hasTools = true
        cleanedContent = ""
      }
    }

    finishReason := "stop"
    var msgContent interface{} = finalContent
    if hasTools {
      finishReason = "tool_calls"
      if cleanedContent == "" {
        msgContent = nil
      } else {
        msgContent = cleanedContent
      }
    }
    w.Header().Set("Content-Type", "application/json")
    return json.NewEncoder(w).Encode(ChatCompletionResponse{
      ID: compID, Object: "chat.completion", Created: created, Model: modelName,
      Choices: []ChatCompletionResponseChoice{{
        Index: 0,
        Message: ChatMessage{
          Role:             "assistant",
          Content:          msgContent,
          ReasoningContent: finalReasoning,
          ToolCalls:        toolCalls,
        },
        FinishReason: finishReason,
      }},
    })
  }

  flusher, _ := w.(http.Flusher)
  streamer := NewStreamer(w, flusher, compID, created, modelName)
  streamer.SendChunk(StreamDelta{Role: "assistant"})

  interceptor := NewStreamToolInterceptor(streamer)
  reader := bufio.NewReader(streamResp.Body)
  var emittedContent, emittedReasoning string

  for {
    line, err := reader.ReadString('\n')
    if err != nil {
      break
    }
    line = strings.TrimSpace(line)
    if strings.HasPrefix(line, "data: ") {
      if c, r, ok := parseGradioSnapshot(strings.TrimPrefix(line, "data: ")); ok {
        if len(r) > len(emittedReasoning) {
          streamer.SendChunk(StreamDelta{ReasoningContent: r[len(emittedReasoning):]})
          emittedReasoning = r
        }
        if len(c) > len(emittedContent) {
          interceptor.ProcessContentDelta(c[len(emittedContent):])
          emittedContent = c
        }
      }
    }
  }

  interceptor.FlushRemaining()
  if !interceptor.HasToolCalls() && len(req.Tools) > 0 {
    if rCalls, _, rHas := DetectToolCalls(emittedReasoning); rHas {
      for _, tc := range rCalls {
        idx := interceptor.toolCallCount
        interceptor.toolCallCount++
        tc.Index = &idx
        streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
      }
    }
  }
  if !interceptor.HasToolCalls() && len(req.Tools) > 0 {
    if directTCs := parseRawToolPayload(emittedContent); len(directTCs) > 0 {
      for _, tc := range directTCs {
        idx := interceptor.toolCallCount
        interceptor.toolCallCount++
        tc.Index = &idx
        streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
      }
    }
  }

  if interceptor.HasToolCalls() {
    streamer.SendChunk(StreamDelta{}, "tool_calls")
  } else {
    streamer.SendChunk(StreamDelta{}, "stop")
  }
  streamer.Done()
  return nil
}

// ---------------------------------------------------------------------------
// HTTP Handlers & Main
// ---------------------------------------------------------------------------

func NewMux(service *KimiService) *http.ServeMux {
  mux := http.NewServeMux()
  mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    if r.Method == http.MethodOptions {
      w.WriteHeader(http.StatusOK)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]interface{}{
      "gateway":          "k3r053n3",
      "status":           "online",
      "model":            service.defaultModel,
      "backend":          service.defaultBackend,
      "reasoning_effort": service.defaultReason,
      "socks5_proxy":     service.socks5Proxy,
      "endpoint":         service.endpoint,
      "routes":           []string{"GET /models", "GET /v1/models", "POST /chat/completions", "POST /v1/chat/completions"},
    })
  })
  modelsH := func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    if r.Method == http.MethodOptions {
      w.WriteHeader(http.StatusOK)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(ModelsResponse{Object: "list", Data: service.ListModels()})
  }
  mux.HandleFunc("/models", modelsH)
  mux.HandleFunc("/v1/models", modelsH)

  chatH := func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    if r.Method == http.MethodOptions {
      w.WriteHeader(http.StatusOK)
      return
    }
    if r.Method != http.MethodPost {
      http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
      return
    }
    var req ChatCompletionRequest
    if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
      http.Error(w, "Invalid JSON payload: "+err.Error(), http.StatusBadRequest)
      return
    }
    if err := service.Chat(w, r, req); err != nil {
      if w.Header().Get("Content-Type") != "text/event-stream" {
        http.Error(w, "Upstream error: "+err.Error(), http.StatusBadGateway)
      }
    }
  }
  mux.HandleFunc("/chat/completions", chatH)
  mux.HandleFunc("/v1/chat/completions", chatH)
  return mux
}

func resolveProxyEnv() string {
  for _, k := range []string{"ALL_PROXY", "all_proxy", "SOCKS5_PROXY", "socks5_proxy", "SOCKS_PROXY", "socks_proxy"} {
    if v := strings.TrimSpace(os.Getenv(k)); v != "" {
      return v
    }
  }
  return ""
}

func main() {
  port := flag.String("port", "8080", "Port to listen on")
  endpoint := flag.String("endpoint", "https://cw-105-kimi-k3-gguf-demo.hf.space", "Root URL of the Kimi K3 Gradio space")
  model := flag.String("model", "kimi-k3", "Exposed default model name")
  backend := flag.String("backend", "direct:together", "Default backend (direct:together, direct:fireworks, hf:fireworks-ai, hf:together, hf:featherless-ai, hf:baseten)")
  reasoning := flag.String("reasoning", "max", "Default reasoning effort (max, high, low, default)")
  tokens := flag.Int("max-tokens", 8192, "Default max completion tokens (256-8192)")
  temp := flag.Float64("temperature", 0.7, "Default temperature (0.0-1.5)")
  socks := flag.String("socks", "", "SOCKS5 proxy URL (e.g. socks5://127.0.0.1:1080 or socks5://user:pass@host:port)")
  proxy := flag.String("proxy", "", "Alias for -socks")
  socks5 := flag.String("socks5", "", "Alias for -socks")
  ua := flag.String("user-agent", DefaultUserAgent, "Custom User-Agent header")
  uaShort := flag.String("ua", "", "Alias for -user-agent")
  flag.Parse()

  ConfiguredUserAgent = *ua
  if *uaShort != "" {
    ConfiguredUserAgent = *uaShort
  }
  socksProxy := *socks
  if socksProxy == "" {
    socksProxy = *proxy
  }
  if socksProxy == "" {
    socksProxy = *socks5
  }
  if socksProxy == "" {
    socksProxy = resolveProxyEnv()
  }

  svc := NewKimiService(*endpoint, *model, *backend, *reasoning, socksProxy, *tokens, *temp)
  fmt.Printf("k3r053n3 starting on port %s...\nTarget Endpoint: %s\nModel Name: %s\nBackend: %s\nReasoning Effort: %s\nMax Tokens: %d\nTemperature: %.2f\n", *port, svc.endpoint, svc.defaultModel, svc.defaultBackend, svc.defaultReason, svc.defaultTokens, svc.defaultTemp)
  if svc.socks5Proxy != "" {
    fmt.Printf("SOCKS5 Proxy: %s\n", svc.socks5Proxy)
  }
  fmt.Printf("Endpoints:\n  GET  http://localhost:%s/v1/models\n  POST http://localhost:%s/v1/chat/completions\n", *port, *port)

  if err := http.ListenAndServe(":"+*port, NewMux(svc)); err != nil {
    fmt.Fprintf(os.Stderr, "server failed: %v\n", err)
    os.Exit(1)
  }
}