// hygate: Standalone OpenAI-compatible gateway for Hunyuan 3 Gradio spaces
// Created by Luxferre in 2026, released into the public domain

package main

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

var (
  DefaultUserAgent    = "Mozilla/5.0 (X11; Linux x86_64; rv:153.0) Gecko/20100101 Firefox/153.0"
  ConfiguredUserAgent string
)

// ---------------------------------------------------------------------------
// 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 {
  if m.Content == nil {
    return ""
  }
  if str, ok := m.Content.(string); ok {
    return str
  }
  if parts, ok := m.Content.([]interface{}); ok {
    var sb strings.Builder
    for _, p := range parts {
      if str, ok := p.(string); ok {
        sb.WriteString(str)
      } else if itemMap, ok := p.(map[string]interface{}); ok {
        if textVal, ok := itemMap["text"].(string); ok {
          sb.WriteString(textVal)
        }
      }
    }
    return sb.String()
  }
  b, err := json.Marshal(m.Content)
  if err == nil {
    return string(b)
  }
  return fmt.Sprintf("%v", m.Content)
}

type ChatCompletionRequest struct {
  Model               string        `json:"model"`
  Messages            []ChatMessage `json:"messages"`
  Tools               []Tool        `json:"tools,omitempty"`
  ToolChoice          interface{}   `json:"tool_choice,omitempty"`
  Stream              bool          `json:"stream"`
  MaxTokens           int           `json:"max_tokens"`
  MaxCompletionTokens int           `json:"max_completion_tokens"`
  Temperature         *float64      `json:"temperature,omitempty"`
  TopP                *float64      `json:"top_p,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"`
}

// ---------------------------------------------------------------------------
// Helper Utilities
// ---------------------------------------------------------------------------

func GenerateUUID() string {
  var b [16]byte
  _, err := rand.Read(b[:])
  if err != nil {
    return "00000000-0000-4000-8000-000000000000"
  }
  b[6] = (b[6] & 0x0f) | 0x40
  b[8] = (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 FibonacciDelay(attempt int) time.Duration {
  if attempt <= 0 {
    return 1 * time.Second
  }
  a, b := 1, 1
  for i := 1; i < attempt; i++ {
    a, b = b, a+b
  }
  return time.Duration(a) * time.Second
}

func DoWithFibonacciRetry(client *http.Client, makeReq func() (*http.Request, error), maxRetries int) (*http.Response, error) {
  var lastErr error
  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 {
      respBody, _ := io.ReadAll(resp.Body)
      resp.Body.Close()
      lastErr = fmt.Errorf("HTTP status %d: %s", resp.StatusCode, string(respBody))
    } else {
      lastErr = err
    }

    if attempt < maxRetries {
      delay := FibonacciDelay(attempt)
      time.Sleep(delay)
    }
  }
  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) int {
  mt := req.MaxTokens
  if mt == 0 && req.MaxCompletionTokens > 0 {
    mt = req.MaxCompletionTokens
  }
  if mt <= 0 {
    mt = 131072
  }
  return mt
}

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
}

// ---------------------------------------------------------------------------
// Response Framing
// ---------------------------------------------------------------------------

type FinalOutput struct {
  Content          interface{}
  ReasoningContent string
  ToolCalls        []ToolCall
  FinishReason     string
}

func WriteCompletionResponse(w http.ResponseWriter, completionID string, created int64, model string, out FinalOutput) {
  finish := out.FinishReason
  if finish == "" {
    finish = "stop"
  }
  resp := ChatCompletionResponse{
    ID:      completionID,
    Object:  "chat.completion",
    Created: created,
    Model:   model,
    Choices: []ChatCompletionResponseChoice{
      {
        Index: 0,
        Message: ChatMessage{
          Role:             "assistant",
          Content:          out.Content,
          ReasoningContent: out.ReasoningContent,
          ToolCalls:        out.ToolCalls,
        },
        FinishReason: finish,
      },
    },
    Usage: Usage{
      PromptTokens:     0,
      CompletionTokens: 0,
      TotalTokens:      0,
    },
  }
  w.Header().Set("Content-Type", "application/json")
  json.NewEncoder(w).Encode(resp)
}

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) Role() {
  sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{Role: "assistant"})
}

func (s *Streamer) Reasoning(text string) {
  sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{ReasoningContent: text})
}

func (s *Streamer) Content(text string) {
  sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{Content: text})
}

func (s *Streamer) ToolCallDelta(tc ToolCall) {
  sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{ToolCalls: []ToolCall{tc}})
}

func (s *Streamer) Finish(reason string) {
  sendStreamChunk(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{}, &reason)
}

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

func sendStreamDelta(w http.ResponseWriter, flusher http.Flusher, completionID string, createdTime int64, modelName string, delta StreamDelta) {
  sendStreamChunk(w, flusher, completionID, createdTime, modelName, delta, nil)
}

func sendStreamChunk(w http.ResponseWriter, flusher http.Flusher, completionID string, createdTime int64, modelName string, delta StreamDelta, finishReason *string) {
  chunk := StreamResponse{
    ID:      completionID,
    Object:  "chat.completion.chunk",
    Created: createdTime,
    Model:   modelName,
    Choices: []StreamChoice{
      {
        Index:        0,
        Delta:        delta,
        FinishReason: finishReason,
      },
    },
  }
  b, _ := json.Marshal(chunk)
  fmt.Fprintf(w, "data: %s\n\n", b)
  if flusher != nil {
    flusher.Flush()
  }
}

// ---------------------------------------------------------------------------
// Hunyuan 3 Provider & Stream Parsing
// ---------------------------------------------------------------------------

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

  inner, ok := raw[0].([]interface{})
  if !ok || len(inner) < 2 {
    return "", "", nil, false
  }

  contentStr, _ := inner[0].(string)
  reasoningStr, _ := inner[1].(string)

  var toolCalls []ToolCall
  if len(inner) >= 3 && inner[2] != nil {
    b, err := json.Marshal(inner[2])
    if err == nil {
      var tcs []ToolCall
      if json.Unmarshal(b, &tcs) == nil && len(tcs) > 0 {
        toolCalls = tcs
      }
    }
  }

  return contentStr, reasoningStr, toolCalls, true
}

type HunyuanService struct {
  endpoint string
  client   *http.Client
}

func NewHunyuanService(endpoint string) *HunyuanService {
  cleanEndpoint := strings.TrimRight(endpoint, "/")
  return &HunyuanService{
    endpoint: cleanEndpoint,
    client:   &http.Client{Timeout: 300 * time.Second},
  }
}

func (s *HunyuanService) ListModels() []ModelItem {
  now := time.Now().Unix()
  return []ModelItem{
    {ID: "hy3", Object: "model", Created: now, OwnedBy: "hunyuan"},
  }
}

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

  var systemPromptStr string
  var historyArray []map[string]interface{}
  var messageStr string

  var nonSystemMsgs []ChatMessage
  for _, msg := range req.Messages {
    cStr := msg.GetContentString()
    if msg.Role == "system" && systemPromptStr == "" {
      systemPromptStr = cStr
    } else {
      nonSystemMsgs = append(nonSystemMsgs, msg)
    }
  }

  if len(nonSystemMsgs) > 0 {
    for i := 0; i < len(nonSystemMsgs)-1; i++ {
      m := nonSystemMsgs[i]
      cStr := m.GetContentString()
      item := map[string]interface{}{"role": m.Role}
      switch m.Role {
      case "assistant":
        if cStr != "" {
          item["content"] = cStr
        } else {
          item["content"] = nil
        }
        if m.ReasoningContent != "" {
          item["reasoning_content"] = m.ReasoningContent
        }
        if len(m.ToolCalls) > 0 {
          item["tool_calls"] = m.ToolCalls
        }
      case "tool", "function":
        item["role"] = "tool"
        item["content"] = cStr
        toolID := m.ToolCallID
        if toolID == "" {
          toolID = m.Name
        }
        if toolID != "" {
          item["tool_call_id"] = toolID
        }
        if m.Name != "" {
          item["name"] = m.Name
        }
      default:
        item["content"] = cStr
      }
      historyArray = append(historyArray, item)
    }

    lastMsg := nonSystemMsgs[len(nonSystemMsgs)-1]
    lastContent := lastMsg.GetContentString()
    if lastMsg.Role == "tool" || lastMsg.Role == "function" {
      toolName := lastMsg.Name
      if toolName == "" {
        toolName = lastMsg.ToolCallID
      }
      if toolName != "" {
        messageStr = fmt.Sprintf("Tool result for %s: %s", toolName, lastContent)
      } else {
        messageStr = lastContent
      }
    } else {
      messageStr = lastContent
    }
  }

  functionsJSONStr := ""
  if len(req.Tools) > 0 {
    b, err := json.Marshal(req.Tools)
    if err == nil {
      functionsJSONStr = string(b)
    }
  }

  var tempVal interface{}
  if req.Temperature != nil {
    tempVal = *req.Temperature
  }
  var topPVal interface{}
  if req.TopP != nil {
    topPVal = *req.TopP
  } else {
    topPVal = 0
  }

  gradioData := []interface{}{
    messageStr,
    systemPromptStr,
    historyArray,
    "high",
    tempVal,
    maxTokens,
    topPVal,
    nil,
    functionsJSONStr,
  }
  gradioPayload := map[string]interface{}{"data": gradioData}
  jsonPayload, err := json.Marshal(gradioPayload)
  if err != nil {
    return fmt.Errorf("failed to encode request: %w", err)
  }

  effUA := EffectiveUserAgent(r)

  callURL := s.endpoint + "/gradio_api/call/chat"
  makeCallReq := func() (*http.Request, error) {
    req, err := http.NewRequest("POST", callURL, bytes.NewBuffer(jsonPayload))
    if err != nil {
      return nil, err
    }
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("User-Agent", effUA)
    return req, nil
  }

  resp, err := DoWithFibonacciRetry(s.client, makeCallReq, 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")
  }

  streamURL := fmt.Sprintf("%s/gradio_api/call/chat/%s", s.endpoint, joinRes.EventID)
  makeStreamReq := func() (*http.Request, error) {
    req, err := http.NewRequest("GET", streamURL, nil)
    if err != nil {
      return nil, err
    }
    req.Header.Set("Accept", "text/event-stream")
    req.Header.Set("User-Agent", effUA)
    return req, nil
  }

  streamResp, err := DoWithFibonacciRetry(s.client, makeStreamReq, 5)
  if err != nil {
    return fmt.Errorf("upstream stream error: %w", err)
  }
  defer streamResp.Body.Close()

  completionID := "chatcmpl-" + GenerateUUID()
  createdTime := time.Now().Unix()

  if !req.Stream {
    reader := bufio.NewReader(streamResp.Body)
    var finalContent, finalReasoning string
    var finalToolCalls []ToolCall

    for {
      line, err := reader.ReadString('\n')
      if err != nil {
        break
      }
      line = strings.TrimSpace(line)
      if strings.HasPrefix(line, "data: ") {
        dataJSON := strings.TrimPrefix(line, "data: ")
        if c, r, tc, ok := parseGradioStreamData(dataJSON); ok {
          finalContent = c
          finalReasoning = r
          if len(tc) > 0 {
            finalToolCalls = tc
          }
        }
      }
    }

    finishReason := "stop"
    var msgContent interface{} = finalContent
    if len(finalToolCalls) > 0 {
      finishReason = "tool_calls"
      if finalContent == "" {
        msgContent = nil
      }
    }

    WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{
      Content:          msgContent,
      ReasoningContent: finalReasoning,
      ToolCalls:        finalToolCalls,
      FinishReason:     finishReason,
    })
    return nil
  }

  flusher, _ := w.(http.Flusher)
  streamer := NewStreamer(w, flusher, completionID, createdTime, modelName)
  streamer.Role()

  reader := bufio.NewReader(streamResp.Body)
  var emittedContent, emittedReasoning string
  lastToolCallArgs := map[string]string{}
  emittedToolCallIDs := map[string]bool{}

  for {
    line, err := reader.ReadString('\n')
    if err != nil {
      break
    }
    line = strings.TrimSpace(line)
    if strings.HasPrefix(line, "data: ") {
      dataJSON := strings.TrimPrefix(line, "data: ")
      if c, r, tc, ok := parseGradioStreamData(dataJSON); ok {
        if len(r) > len(emittedReasoning) {
          rDelta := r[len(emittedReasoning):]
          emittedReasoning = r
          streamer.Reasoning(rDelta)
        }
        if len(c) > len(emittedContent) {
          cDelta := c[len(emittedContent):]
          emittedContent = c
          streamer.Content(cDelta)
        }
        if len(tc) > 0 {
          for idx, t := range tc {
            key := t.ID
            if key == "" {
              key = fmt.Sprintf("idx_%d", idx)
            }
            prevArgs := lastToolCallArgs[key]
            currArgs := t.Function.Arguments
            if !emittedToolCallIDs[key] {
              emittedToolCallIDs[key] = true
              lastToolCallArgs[key] = currArgs
              tCopy := ToolCall{
                Index:    new(int),
                ID:       t.ID,
                Type:     "function",
                Function: ToolCallFunction{Name: t.Function.Name, Arguments: currArgs},
              }
              *tCopy.Index = idx
              streamer.ToolCallDelta(tCopy)
            } else if len(currArgs) > len(prevArgs) {
              argDelta := currArgs[len(prevArgs):]
              lastToolCallArgs[key] = currArgs
              tCopy := ToolCall{
                Index:    new(int),
                Function: ToolCallFunction{Arguments: argDelta},
              }
              *tCopy.Index = idx
              streamer.ToolCallDelta(tCopy)
            }
          }
        }
      }
    }
  }

  if len(emittedToolCallIDs) > 0 {
    streamer.Finish("tool_calls")
  } else {
    streamer.Finish("stop")
  }
  streamer.Done()
  return nil
}

// ---------------------------------------------------------------------------
// HTTP Server Wiring & Main Entry Point
// ---------------------------------------------------------------------------

func ModelsHandler(service *HunyuanService) http.HandlerFunc {
  return func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    if r.Method == http.MethodOptions {
      w.WriteHeader(http.StatusOK)
      return
    }
    if r.Method != http.MethodGet {
      http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
      return
    }
    items := service.ListModels()
    res := ModelsResponse{Object: "list", Data: items}
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(res)
  }
}

func ChatHandler(service *HunyuanService) http.HandlerFunc {
  return 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)
      }
    }
  }
}

func NewMux(service *HunyuanService) *http.ServeMux {
  mux := http.NewServeMux()
  mh := ModelsHandler(service)
  ch := ChatHandler(service)
  mux.HandleFunc("/models", mh)
  mux.HandleFunc("/v1/models", mh)
  mux.HandleFunc("/chat/completions", ch)
  mux.HandleFunc("/v1/chat/completions", ch)
  return mux
}

func main() {
  portFlag := flag.String("port", "8080", "Port to listen on")
  endpointFlag := flag.String("endpoint", "https://tencent-hy3.hf.space", "Root URL of the Hunyuan Gradio space")
  uaFlag := flag.String("user-agent", DefaultUserAgent, "Custom User-Agent header")
  uaShortFlag := flag.String("ua", "", "Alias for -user-agent")
  flag.Parse()

  ConfiguredUserAgent = *uaFlag
  if *uaShortFlag != "" {
    ConfiguredUserAgent = *uaShortFlag
  }

  service := NewHunyuanService(*endpointFlag)
  mux := NewMux(service)

  fmt.Printf("hygate starting on port %s...\n", *portFlag)
  fmt.Printf("Target Endpoint: %s\n", service.endpoint)
  fmt.Printf("User-Agent: %s\n", ConfiguredUserAgent)
  fmt.Printf("Endpoints:\n")
  fmt.Printf("  GET  http://localhost:%s/v1/models\n", *portFlag)
  fmt.Printf("  POST http://localhost:%s/v1/chat/completions\n", *portFlag)

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