qorona/main.go

// Qorona: Standalone OpenAI-compatible gateway for Qwen 3.8 Max HuggingFace Spaces
// Created by Luxferre in 2026, released into the public domain

package main

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

var (
  DefaultSpaceURL     = "https://harpreetsahota-qwen38-max-openlogo-demo.hf.space"
  DefaultSamplePath   = "/home/user/datasets/openlogo/data/data_0/logos32plus_002359.jpg"
  DefaultModelID      = "qwen-3.8-max"
  DefaultUserAgent    = "Mozilla/5.0 (X11; Linux x86_64; rv:153.0) Gecko/20100101 Firefox/153.0"
  ConfiguredUserAgent string
  ConfiguredToken     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 StreamOptions struct {
  IncludeUsage bool `json:"include_usage,omitempty"`
}

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"`
  StreamOptions       *StreamOptions `json:"stream_options,omitempty"`
  MaxTokens           int            `json:"max_tokens"`
  MaxCompletionTokens int            `json:"max_completion_tokens"`
  Temperature         *float64       `json:"temperature,omitempty"`
  TopP                *float64       `json:"top_p,omitempty"`
  ReasoningEffort     string         `json:"reasoning_effort,omitempty"`
  Thinking            interface{}    `json:"thinking,omitempty"`
}

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

type PromptTokensDetails struct {
  CachedTokens int `json:"cached_tokens"`
}

type CompletionTokensDetails struct {
  ReasoningTokens int `json:"reasoning_tokens,omitempty"`
}

type Usage struct {
  PromptTokens            int                      `json:"prompt_tokens"`
  CompletionTokens        int                      `json:"completion_tokens"`
  TotalTokens             int                      `json:"total_tokens"`
  CachedTokens            int                      `json:"cached_tokens,omitempty"`
  PromptTokensDetails     *PromptTokensDetails     `json:"prompt_tokens_details,omitempty"`
  CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
}

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"`
  Usage   *Usage         `json:"usage,omitempty"`
}

// ---------------------------------------------------------------------------
// 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 || resp.StatusCode == http.StatusCreated || resp.StatusCode == http.StatusNoContent) {
      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, X-HF-Token")
}

func ResolveMaxTokens(req ChatCompletionRequest) int {
  mt := req.MaxTokens
  if mt == 0 && req.MaxCompletionTokens > 0 {
    mt = req.MaxCompletionTokens
  }
  if mt <= 0 {
    mt = 4000
  }
  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
}

func EffectiveToken(r *http.Request) string {
  if r != nil {
    if hf := r.Header.Get("X-HF-Token"); hf != "" {
      return hf
    }
    if auth := r.Header.Get("Authorization"); auth != "" {
      if strings.HasPrefix(strings.ToLower(auth), "bearer ") {
        token := strings.TrimSpace(auth[7:])
        if token != "" && token != "-" {
          return token
        }
      }
    }
  }
  if ConfiguredToken != "" {
    return ConfiguredToken
  }
  if envTok := os.Getenv("HF_TOKEN"); envTok != "" {
    return envTok
  }
  return ""
}

// ---------------------------------------------------------------------------
// FiftyOne HF Space Session Manager
// ---------------------------------------------------------------------------

type FiftyOneSessionManager struct {
  spaceURL    string
  client      *http.Client
  mu          sync.Mutex
  token       string
  datasetName string
  lastSeen    time.Time
}

func NewFiftyOneSessionManager(spaceURL string, client *http.Client) *FiftyOneSessionManager {
  sm := &FiftyOneSessionManager{
    spaceURL: strings.TrimRight(spaceURL, "/"),
    client:   client,
  }
  go sm.heartbeatLoop()
  return sm
}

func (sm *FiftyOneSessionManager) EnsureSession(ctx context.Context, reqUserAgent string, hfToken string) (string, string, error) {
  sm.mu.Lock()
  defer sm.mu.Unlock()

  if sm.token != "" && sm.datasetName != "" && time.Since(sm.lastSeen) < 20*time.Minute {
    sm.lastSeen = time.Now()
    return sm.token, sm.datasetName, nil
  }

  return sm.startNewSessionLocked(ctx, reqUserAgent, hfToken)
}

func (sm *FiftyOneSessionManager) startNewSessionLocked(ctx context.Context, reqUserAgent string, hfToken string) (string, string, error) {
  url := sm.spaceURL + "/__session/start"
  req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader([]byte("{}")))
  if err != nil {
    return "", "", err
  }
  req.Header.Set("Content-Type", "application/json")
  if reqUserAgent != "" {
    req.Header.Set("User-Agent", reqUserAgent)
  } else {
    req.Header.Set("User-Agent", DefaultUserAgent)
  }
  if hfToken != "" {
    req.Header.Set("Authorization", "Bearer "+hfToken)
  }

  resp, err := sm.client.Do(req)
  if err != nil {
    return "", "", fmt.Errorf("failed to connect to session start: %w", err)
  }
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    bodyBytes, _ := io.ReadAll(resp.Body)
    return "", "", fmt.Errorf("session start returned status %d: %s", resp.StatusCode, string(bodyBytes))
  }

  var startRes struct {
    Token string `json:"token"`
    URL   string `json:"url"`
    Error string `json:"error"`
  }
  if err := json.NewDecoder(resp.Body).Decode(&startRes); err != nil {
    return "", "", fmt.Errorf("failed to decode session start json: %w", err)
  }

  if startRes.Error != "" {
    return "", "", fmt.Errorf("session start error: %s", startRes.Error)
  }
  if startRes.Token == "" {
    return "", "", fmt.Errorf("session start returned empty token")
  }

  parts := strings.Split(strings.TrimRight(startRes.URL, "/"), "/")
  datasetName := parts[len(parts)-1]
  if datasetName == "" {
    datasetName = "openlogo-session-" + startRes.Token
  }

  sm.token = startRes.Token
  sm.datasetName = datasetName
  sm.lastSeen = time.Now()

  log.Printf("[Session] Allocated new session token=%s dataset=%s", sm.token, sm.datasetName)
  return sm.token, sm.datasetName, nil
}

func (sm *FiftyOneSessionManager) Invalidate() {
  sm.mu.Lock()
  defer sm.mu.Unlock()
  sm.token = ""
  sm.datasetName = ""
}

func (sm *FiftyOneSessionManager) heartbeatLoop() {
  ticker := time.NewTicker(45 * time.Second)
  defer ticker.Stop()

  for range ticker.C {
    sm.mu.Lock()
    token := sm.token
    dataset := sm.datasetName
    sm.mu.Unlock()

    if token == "" || dataset == "" {
      continue
    }

    url := fmt.Sprintf("%s/__session/heartbeat?dataset=%s", sm.spaceURL, dataset)
    req, err := http.NewRequest("POST", url, nil)
    if err == nil {
      req.Header.Set("X-FiftyOne-Session", token)
      req.Header.Set("User-Agent", DefaultUserAgent)
      if ConfiguredToken != "" {
        req.Header.Set("Authorization", "Bearer "+ConfiguredToken)
      }
      resp, err := sm.client.Do(req)
      if err == nil {
        resp.Body.Close()
        if resp.StatusCode == http.StatusUnauthorized {
          sm.Invalidate()
        }
      }
    }
  }
}

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

type FinalOutput struct {
  Content          interface{}
  ReasoningContent string
  ToolCalls        []ToolCall
  FinishReason     string
  PromptTokens     int
  CompletionTokens int
  CachedTokens     int
  ReasoningTokens  int
}

func EstimateCachedTokens(promptTokens int, currentQuestion string) int {
  if promptTokens <= 0 {
    return 0
  }
  qLen := len(strings.TrimSpace(currentQuestion))
  uncached := (qLen + 3) / 4
  if uncached < 1 {
    uncached = 1
  }
  if uncached >= promptTokens {
    return 0
  }
  return promptTokens - uncached
}

func ResolveCachedTokens(upstreamCached int, promptTokens int, currentQuestion string) int {
  if upstreamCached > 0 {
    return upstreamCached
  }
  return EstimateCachedTokens(promptTokens, currentQuestion)
}

func BuildUsage(promptTokens, completionTokens, cachedTokens, reasoningTokens int) Usage {
  totalTokens := promptTokens + completionTokens
  return Usage{
    PromptTokens:     promptTokens,
    CompletionTokens: completionTokens,
    TotalTokens:      totalTokens,
    CachedTokens:     cachedTokens,
    PromptTokensDetails: &PromptTokensDetails{
      CachedTokens: cachedTokens,
    },
    CompletionTokensDetails: &CompletionTokensDetails{
      ReasoningTokens: reasoningTokens,
    },
  }
}

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: BuildUsage(out.PromptTokens, out.CompletionTokens, out.CachedTokens, out.ReasoningTokens),
  }
  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) EmitUsage(usage Usage) {
  chunk := StreamResponse{
    ID:      s.id,
    Object:  "chat.completion.chunk",
    Created: s.created,
    Model:   s.model,
    Choices: []StreamChoice{},
    Usage:   &usage,
  }
  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 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()
  }
}

// ---------------------------------------------------------------------------
// Tool Calling and Reasoning Extraction
// ---------------------------------------------------------------------------

func cleanJSONBlock(input string) string {
  s := strings.TrimSpace(input)
  if strings.HasPrefix(s, "```") {
    lines := strings.Split(s, "\n")
    if len(lines) >= 2 {
      if strings.HasPrefix(lines[len(lines)-1], "```") {
        lines = lines[1 : len(lines)-1]
      } else {
        lines = lines[1:]
      }
      s = strings.TrimSpace(strings.Join(lines, "\n"))
    }
  }
  return s
}

func sanitizeJSONValue(v interface{}) interface{} {
  switch val := v.(type) {
  case string:
    return strings.TrimSpace(val)
  case map[string]interface{}:
    cleanMap := make(map[string]interface{})
    for k, childV := range val {
      cleanKey := strings.TrimSpace(k)
      cleanMap[cleanKey] = sanitizeJSONValue(childV)
    }
    return cleanMap
  case []interface{}:
    cleanSlice := make([]interface{}, len(val))
    for i, childV := range val {
      cleanSlice[i] = sanitizeJSONValue(childV)
    }
    return cleanSlice
  default:
    return v
  }
}

func parseSingleToolCall(jsonStr string) (ToolCall, bool) {
  cleaned := cleanJSONBlock(jsonStr)
  var raw map[string]interface{}
  if err := json.Unmarshal([]byte(cleaned), &raw); err == nil {
    sanitizedRaw, ok := sanitizeJSONValue(raw).(map[string]interface{})
    if !ok {
      sanitizedRaw = raw
    }

    for _, wrapperKey := range []string{"function", "function_call", "tool_call"} {
      if fnObj, ok := sanitizedRaw[wrapperKey].(map[string]interface{}); ok {
        if nameVal, ok := fnObj["name"].(string); ok && nameVal != "" {
          argsStr := "{}"
          var argsVal interface{}
          if a, hasA := fnObj["arguments"]; hasA {
            argsVal = a
          } else if p, hasP := fnObj["parameters"]; hasP {
            argsVal = p
          } else if args, hasArgs := fnObj["args"]; hasArgs {
            argsVal = args
          }
          if argsVal != nil {
            if s, isStr := argsVal.(string); isStr {
              var innerObj interface{}
              if json.Unmarshal([]byte(s), &innerObj) == nil {
                b, _ := json.Marshal(sanitizeJSONValue(innerObj))
                argsStr = string(b)
              } else {
                argsStr = strings.TrimSpace(s)
              }
            } else {
              b, _ := json.Marshal(sanitizeJSONValue(argsVal))
              argsStr = string(b)
            }
          }
          return ToolCall{
            ID:   "call_" + GenerateUUID()[:8],
            Type: "function",
            Function: ToolCallFunction{
              Name:      nameVal,
              Arguments: argsStr,
            },
          }, true
        }
      }
    }

    nameVal := ""
    for _, key := range []string{"name", "function", "action", "call"} {
      if n, ok := sanitizedRaw[key].(string); ok && n != "" {
        nameVal = n
        break
      }
    }

    if nameVal != "" {
      argsStr := "{}"
      var argsVal interface{}
      for _, key := range []string{"arguments", "parameters", "args", "input"} {
        if a, ok := sanitizedRaw[key]; ok {
          argsVal = a
          break
        }
      }
      if argsVal != nil {
        if s, isStr := argsVal.(string); isStr {
          var innerObj interface{}
          if json.Unmarshal([]byte(s), &innerObj) == nil {
            b, _ := json.Marshal(sanitizeJSONValue(innerObj))
            argsStr = string(b)
          } else {
            argsStr = strings.TrimSpace(s)
          }
        } else {
          b, _ := json.Marshal(sanitizeJSONValue(argsVal))
          argsStr = string(b)
        }
      } else {
        argsMap := make(map[string]interface{})
        for k, v := range sanitizedRaw {
          if k != "name" && k != "function" && k != "type" && k != "action" && k != "call" {
            argsMap[k] = v
          }
        }
        if len(argsMap) > 0 {
          b, _ := json.Marshal(sanitizeJSONValue(argsMap))
          argsStr = string(b)
        }
      }
      return ToolCall{
        ID:   "call_" + GenerateUUID()[:8],
        Type: "function",
        Function: ToolCallFunction{
          Name:      nameVal,
          Arguments: argsStr,
        },
      }, true
    }
  }

  return ToolCall{}, false
}

func parseXMLToolCall(block string) (ToolCall, bool) {
  inner := strings.TrimSpace(block)
  if strings.HasPrefix(inner, "<tool_call>") {
    inner = strings.TrimPrefix(inner, "<tool_call>")
  }
  if strings.HasSuffix(inner, "</tool_call>") {
    inner = strings.TrimSuffix(inner, "</tool_call>")
  }
  inner = cleanJSONBlock(inner)

  if tc, ok := parseSingleToolCall(inner); ok {
    return tc, true
  }

  var fnName string
  if strings.Contains(inner, "<name>") && strings.Contains(inner, "</name>") {
    nStart := strings.Index(inner, "<name>") + len("<name>")
    nEnd := strings.Index(inner, "</name>")
    if nStart < nEnd {
      fnName = strings.TrimSpace(inner[nStart:nEnd])
    }
  } else if strings.Contains(inner, "<function_name>") && strings.Contains(inner, "</function_name>") {
    nStart := strings.Index(inner, "<function_name>") + len("<function_name>")
    nEnd := strings.Index(inner, "</function_name>")
    if nStart < nEnd {
      fnName = strings.TrimSpace(inner[nStart:nEnd])
    }
  }

  var argsStr string
  if strings.Contains(inner, "<arguments>") && strings.Contains(inner, "</arguments>") {
    aStart := strings.Index(inner, "<arguments>") + len("<arguments>")
    aEnd := strings.Index(inner, "</arguments>")
    if aStart < aEnd {
      argsStr = strings.TrimSpace(inner[aStart:aEnd])
    }
  }

  if fnName != "" {
    if argsStr == "" {
      argsStr = "{}"
    }
    return ToolCall{
      ID:   "call_" + GenerateUUID()[:8],
      Type: "function",
      Function: ToolCallFunction{
        Name:      fnName,
        Arguments: argsStr,
      },
    }, true
  }

  return ToolCall{}, false
}

func ExtractToolCallBlocks(content string) (blocks []string, remaining string) {
  s := content
  remaining = content

  for strings.Contains(s, "<tool_call>") {
    sIdx := strings.Index(s, "<tool_call>")
    rest := s[sIdx+len("<tool_call>"):]

    relNextSIdx := strings.Index(rest, "<tool_call>")
    var nextSIdx int
    if relNextSIdx != -1 {
      nextSIdx = sIdx + len("<tool_call>") + relNextSIdx
    } else {
      nextSIdx = -1
    }

    relEIdx := strings.Index(rest, "</tool_call>")
    var eIdx int
    if relEIdx != -1 {
      eIdx = sIdx + len("<tool_call>") + relEIdx
    } else {
      eIdx = -1
    }

    var blockText string
    var blockEndPos int

    if eIdx != -1 && (nextSIdx == -1 || eIdx < nextSIdx) {
      blockEndPos = eIdx + len("</tool_call>")
      blockText = s[sIdx:blockEndPos]
      s = s[blockEndPos:]
    } else if nextSIdx != -1 {
      blockEndPos = nextSIdx
      blockText = s[sIdx:blockEndPos]
      s = s[blockEndPos:]
    } else {
      blockText = s[sIdx:]
      s = ""
    }

    blocks = append(blocks, blockText)
  }

  for strings.Contains(remaining, "<tool_call>") {
    st := strings.Index(remaining, "<tool_call>")
    rest := remaining[st+len("<tool_call>"):]

    relNext := strings.Index(rest, "<tool_call>")
    var nextSt int
    if relNext != -1 {
      nextSt = st + len("<tool_call>") + relNext
    } else {
      nextSt = -1
    }

    relEn := strings.Index(rest, "</tool_call>")
    var en int
    if relEn != -1 {
      en = st + len("<tool_call>") + relEn
    } else {
      en = -1
    }

    if en != -1 && (nextSt == -1 || en < nextSt) {
      remaining = strings.TrimSpace(remaining[:st] + remaining[en+len("</tool_call>"):])
    } else if nextSt != -1 {
      remaining = strings.TrimSpace(remaining[:st] + remaining[nextSt:])
    } else {
      remaining = strings.TrimSpace(remaining[:st])
    }
  }

  return blocks, remaining
}

func DetectToolCalls(content string) (calls []ToolCall, remainingText string) {
  trimmed := strings.TrimSpace(content)
  if trimmed == "" {
    return nil, ""
  }

  if strings.Contains(content, "<tool_call>") {
    blocks, rem := ExtractToolCallBlocks(content)
    for _, blk := range blocks {
      if tc, ok := parseXMLToolCall(blk); ok {
        calls = append(calls, tc)
      }
    }
    if len(calls) > 0 {
      return calls, rem
    }
  }

  var rawArray []interface{}
  cleanedContent := cleanJSONBlock(trimmed)
  if strings.HasPrefix(cleanedContent, "[") && strings.HasSuffix(cleanedContent, "]") {
    if err := json.Unmarshal([]byte(cleanedContent), &rawArray); err == nil && len(rawArray) > 0 {
      allValid := true
      var tempCalls []ToolCall
      for _, item := range rawArray {
        b, _ := json.Marshal(item)
        if tc, ok := parseSingleToolCall(string(b)); ok {
          tempCalls = append(tempCalls, tc)
        } else {
          allValid = false
          break
        }
      }
      if allValid && len(tempCalls) > 0 {
        return tempCalls, ""
      }
    }
  }

  if tc, ok := parseSingleToolCall(trimmed); ok {
    return []ToolCall{tc}, ""
  }

  return nil, content
}

// ---------------------------------------------------------------------------
// Stateful Streaming Tool Call Filter
// ---------------------------------------------------------------------------

type StreamToolCallFilter struct {
  streamer       *Streamer
  buffer         string
  insideToolCall bool
  insideJSON     bool
  hasEmittedTool bool
}

func NewStreamToolCallFilter(streamer *Streamer) *StreamToolCallFilter {
  return &StreamToolCallFilter{
    streamer: streamer,
  }
}

func (f *StreamToolCallFilter) Feed(chunk string) {
  f.buffer += chunk

  for {
    if !f.insideToolCall && !f.insideJSON {
      idx := strings.Index(f.buffer, "<tool_call>")
      if idx != -1 {
        cleanPrefix := f.buffer[:idx]
        if cleanPrefix != "" {
          f.streamer.Content(cleanPrefix)
        }
        f.buffer = f.buffer[idx:]
        f.insideToolCall = true
        continue
      }

      trimmedBuf := strings.TrimSpace(f.buffer)
      if (strings.HasPrefix(trimmedBuf, "{\"name\"") || strings.HasPrefix(trimmedBuf, "{\"function\"")) && strings.HasSuffix(trimmedBuf, "}") {
        if tc, ok := parseSingleToolCall(trimmedBuf); ok {
          f.hasEmittedTool = true
          f.streamer.ToolCallDelta(tc)
          f.buffer = ""
          return
        }
      }

      const safeLen = 12
      if len(f.buffer) > safeLen {
        emitLen := len(f.buffer) - safeLen
        f.streamer.Content(f.buffer[:emitLen])
        f.buffer = f.buffer[emitLen:]
      }
      return
    }

    if f.insideToolCall {
      eIdx := strings.Index(f.buffer, "</tool_call>")
      if eIdx != -1 {
        block := f.buffer[:eIdx+len("</tool_call>")]
        f.buffer = f.buffer[eIdx+len("</tool_call>"):]
        f.insideToolCall = false

        if tc, ok := parseXMLToolCall(block); ok {
          f.hasEmittedTool = true
          f.streamer.ToolCallDelta(tc)
        }
        continue
      }
      return
    }

    return
  }
}

func (f *StreamToolCallFilter) Flush() {
  if f.buffer == "" {
    return
  }
  if f.insideToolCall {
    if tc, ok := parseXMLToolCall(f.buffer); ok {
      f.hasEmittedTool = true
      f.streamer.ToolCallDelta(tc)
      f.buffer = ""
      return
    }
  }

  if tc, ok := parseSingleToolCall(f.buffer); ok {
    f.hasEmittedTool = true
    f.streamer.ToolCallDelta(tc)
    f.buffer = ""
    return
  }

  f.streamer.Content(f.buffer)
  f.buffer = ""
}

func (f *StreamToolCallFilter) HasEmittedTools() bool {
  return f.hasEmittedTool
}

// ---------------------------------------------------------------------------
// Thinking Filter
// ---------------------------------------------------------------------------

var thinkTagRegex = regexp.MustCompile(`(?s)<think>(.*?)(?:</think>|$)`)

func ExtractThinkingContent(text string) (thinking string, cleanText string) {
  matches := thinkTagRegex.FindAllStringSubmatch(text, -1)
  if len(matches) > 0 {
    var thinkParts []string
    for _, m := range matches {
      if len(m) > 1 {
        thinkParts = append(thinkParts, strings.TrimSpace(m[1]))
      }
    }
    thinking = strings.Join(thinkParts, "\n\n")
    cleanText = thinkTagRegex.ReplaceAllString(text, "")
    cleanText = strings.TrimSpace(cleanText)
    return thinking, cleanText
  }
  return "", text
}

// ---------------------------------------------------------------------------
// Conversation Message Formatting
// ---------------------------------------------------------------------------

func formatToolPrompt(tools []Tool) string {
  if len(tools) == 0 {
    return ""
  }
  var sb strings.Builder
  sb.WriteString("\n\n[AVAILABLE TOOLS]\n")
  sb.WriteString("You have access to the following tools/functions:\n")
  for _, t := range tools {
    b, err := json.Marshal(t)
    if err == nil {
      sb.WriteString(string(b))
      sb.WriteString("\n")
    }
  }
  sb.WriteString("\nTo invoke a tool, output a single tool call formatted exactly like this:\n")
  sb.WriteString("<tool_call>{\"name\": \"function_name\", \"arguments\": {\"param1\": \"val1\"}}</tool_call>\n")
  sb.WriteString("Do not add extra explanation when calling a tool.\n\n")
  return sb.String()
}

func PrepareConversation(req ChatCompletionRequest) (history []map[string]interface{}, currentQuestion string, thinkingMode string) {
  var systemParts []string

  // Add default baseline directive to maintain conversational focus and ignore demo image
  systemParts = append(systemParts, "CRITICAL DIRECTIVE: You are operating strictly as a general-purpose text, coding, and reasoning assistant. Completely disregard and ignore any attached background images and image dimension hints unless the user explicitly asks about them. Focus exclusively on the user's instructions, queries, and tool invocations.")

  if len(req.Tools) > 0 {
    systemParts = append(systemParts, formatToolPrompt(req.Tools))
  }

  // Determine thinking mode
  thinkingMode = "auto"
  if req.Thinking != nil {
    switch tv := req.Thinking.(type) {
    case bool:
      if tv {
        thinkingMode = "true"
      } else {
        thinkingMode = "false"
      }
    case string:
      s := strings.ToLower(strings.TrimSpace(tv))
      if s == "true" || s == "on" || s == "enabled" {
        thinkingMode = "true"
      } else if s == "false" || s == "off" || s == "disabled" {
        thinkingMode = "false"
      }
    }
  } else if req.ReasoningEffort != "" {
    re := strings.ToLower(strings.TrimSpace(req.ReasoningEffort))
    if re == "none" || re == "false" || re == "off" {
      thinkingMode = "false"
    } else if re == "low" || re == "medium" || re == "high" || re == "true" {
      thinkingMode = "true"
    }
  }

  var turns []map[string]string
  for _, m := range req.Messages {
    role := strings.ToLower(strings.TrimSpace(m.Role))
    content := m.GetContentString()

    if role == "system" {
      if content != "" {
        systemParts = append(systemParts, content)
      }
      continue
    }

    if role == "tool" || role == "function" {
      toolName := m.Name
      if toolName == "" {
        toolName = "tool"
      }
      content = fmt.Sprintf("<tool_response>{\"name\": %q, \"content\": %s}</tool_response>", toolName, content)
      role = "user"
    }

    if len(m.ToolCalls) > 0 {
      var tcParts []string
      for _, tc := range m.ToolCalls {
        tcParts = append(tcParts, fmt.Sprintf("<tool_call>{\"name\": %q, \"arguments\": %s}</tool_call>", tc.Function.Name, tc.Function.Arguments))
      }
      if content != "" {
        content = content + "\n" + strings.Join(tcParts, "\n")
      } else {
        content = strings.Join(tcParts, "\n")
      }
    }

    turns = append(turns, map[string]string{
      "role":    role,
      "content": content,
    })
  }

  systemText := strings.TrimSpace(strings.Join(systemParts, "\n\n"))

  if len(turns) == 0 {
    emptyHist := make([]map[string]interface{}, 0)
    if systemText != "" {
      return emptyHist, systemText, thinkingMode
    }
    return emptyHist, "Hello", thinkingMode
  }

  lastTurn := turns[len(turns)-1]
  priorTurns := turns[:len(turns)-1]

  historyItems := make([]map[string]interface{}, 0)
  for idx, t := range priorTurns {
    c := t["content"]
    if idx == 0 && systemText != "" {
      c = "[SYSTEM INSTRUCTION]\n" + systemText + "\n\n" + c
    }
    historyItems = append(historyItems, map[string]interface{}{
      "role":    t["role"],
      "content": c,
    })
  }

  currentQuestion = lastTurn["content"]
  if len(priorTurns) == 0 && systemText != "" {
    currentQuestion = "[SYSTEM INSTRUCTION]\n" + systemText + "\n\n" + currentQuestion
  }

  return historyItems, currentQuestion, thinkingMode
}

// ---------------------------------------------------------------------------
// Main Gateway Service
// ---------------------------------------------------------------------------

type QoronaGateway struct {
  spaceURL   string
  samplePath string
  modelID    string
  client     *http.Client
  sessionMgr *FiftyOneSessionManager
}

func NewQoronaGateway(spaceURL string, samplePath string, modelID string, timeout time.Duration) *QoronaGateway {
  tr := &http.Transport{
    MaxIdleConns:        100,
    MaxIdleConnsPerHost: 20,
    IdleConnTimeout:     90 * time.Second,
  }
  client := &http.Client{
    Transport: tr,
    Timeout:   timeout,
  }
  return &QoronaGateway{
    spaceURL:   strings.TrimRight(spaceURL, "/"),
    samplePath: samplePath,
    modelID:    modelID,
    client:     client,
    sessionMgr: NewFiftyOneSessionManager(spaceURL, client),
  }
}

func (gw *QoronaGateway) HandleModels(w http.ResponseWriter, r *http.Request) {
  EnableCORS(w)
  if r.Method == http.MethodOptions {
    w.WriteHeader(http.StatusNoContent)
    return
  }

  now := time.Now().Unix()
  resp := ModelsResponse{
    Object: "list",
    Data: []ModelItem{
      {
        ID:      gw.modelID,
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
      {
        ID:      "qorona",
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
      {
        ID:      "qwen-3.8-max",
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
      {
        ID:      "qwen3.8-max",
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
      {
        ID:      "qwen38-max",
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
      {
        ID:      "Qwen/Qwen3.8-Max",
        Object:  "model",
        Created: now,
        OwnedBy: "qwen",
      },
    },
  }

  w.Header().Set("Content-Type", "application/json")
  json.NewEncoder(w).Encode(resp)
}

func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Request) {
  EnableCORS(w)
  if r.Method == http.MethodOptions {
    w.WriteHeader(http.StatusNoContent)
    return
  }
  if r.Method != http.MethodPost {
    http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
    return
  }

  var req ChatCompletionRequest
  bodyBytes, err := io.ReadAll(r.Body)
  if err != nil {
    http.Error(w, "Failed to read request body", http.StatusBadRequest)
    return
  }

  if err := json.Unmarshal(bodyBytes, &req); err != nil {
    http.Error(w, fmt.Sprintf("Invalid JSON request: %v", err), http.StatusBadRequest)
    return
  }

  modelName := req.Model
  if modelName == "" {
    modelName = gw.modelID
  }

  maxTokens := ResolveMaxTokens(req)
  history, question, thinkingMode := PrepareConversation(req)
  if history == nil {
    history = make([]map[string]interface{}, 0)
  }

  reqUA := EffectiveUserAgent(r)
  reqHF := EffectiveToken(r)

  // Obtain or refresh FiftyOne Session
  token, datasetName, err := gw.sessionMgr.EnsureSession(r.Context(), reqUA, reqHF)
  if err != nil {
    log.Printf("[Error] Failed to obtain FiftyOne session: %v", err)
    http.Error(w, fmt.Sprintf("Upstream session error: %v", err), http.StatusBadGateway)
    return
  }

  // Execute ask operator
  askPayload := map[string]interface{}{
    "operator_uri":     "@harpreetsahota/qwen38-max/qwen38_chat",
    "dataset_name":     datasetName,
    "view":             []interface{}{},
    "selected":         []interface{}{},
    "selected_samples": []interface{}{},
    "selected_labels":  []interface{}{},
    "filters":          map[string]interface{}{},
    "extended":         map[string]interface{}{},
    "params": map[string]interface{}{
      "panel_id":       "qwen38_chat",
      "__method__":     "ask",
      "filepath":       gw.samplePath,
      "question":       question,
      "history":        history,
      "thinking":       thinkingMode,
      "hint_format":    "auto",
      "hint_text":      "",
      "image_max_side": 1280,
      "max_tokens":     maxTokens,
    },
  }

  askBody, _ := json.Marshal(askPayload)
  executeURL := gw.spaceURL + "/operators/execute"

  makeAskReq := func() (*http.Request, error) {
    hReq, err := http.NewRequestWithContext(r.Context(), "POST", executeURL, bytes.NewReader(askBody))
    if err != nil {
      return nil, err
    }
    hReq.Header.Set("Content-Type", "application/json")
    hReq.Header.Set("X-FiftyOne-Session", token)
    hReq.Header.Set("User-Agent", reqUA)
    if reqHF != "" {
      hReq.Header.Set("Authorization", "Bearer "+reqHF)
    }
    return hReq, nil
  }

  askResp, err := DoWithFibonacciRetry(gw.client, makeAskReq, 3)
  if err != nil {
    log.Printf("[Error] ask execution failed: %v", err)
    gw.sessionMgr.Invalidate()
    http.Error(w, fmt.Sprintf("Upstream operator execution failed: %v", err), http.StatusBadGateway)
    return
  }
  defer askResp.Body.Close()

  var askResult struct {
    Result struct {
      Status string `json:"status"`
      RunID  string `json:"run_id"`
      Error  string `json:"error"`
    } `json:"result"`
    Error        string `json:"error"`
    ErrorMessage string `json:"error_message"`
  }

  if err := json.NewDecoder(askResp.Body).Decode(&askResult); err != nil {
    http.Error(w, fmt.Sprintf("Failed to parse operator response: %v", err), http.StatusBadGateway)
    return
  }

  if askResult.Error != "" || askResult.ErrorMessage != "" {
    errMsg := askResult.Error
    if errMsg == "" {
      errMsg = askResult.ErrorMessage
    }
    http.Error(w, fmt.Sprintf("Upstream error: %s", errMsg), http.StatusBadGateway)
    return
  }

  runID := askResult.Result.RunID
  if runID == "" {
    http.Error(w, "Upstream did not return a run_id", http.StatusBadGateway)
    return
  }

  log.Printf("[Inference] Started run_id=%s model=%s stream=%t thinking=%s", runID, modelName, req.Stream, thinkingMode)

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

  if req.Stream {
    flusher, ok := w.(http.Flusher)
    if !ok {
      http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
      return
    }

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

    toolFilter := NewStreamToolCallFilter(streamer)

    streamCursor := 0
    thinkingCursor := 0
    pollTicker := time.NewTicker(150 * time.Millisecond)
    defer pollTicker.Stop()

    for {
      select {
      case <-r.Context().Done():
        log.Printf("[Stream] Client disconnected for run_id=%s", runID)
        return
      case <-pollTicker.C:
        // 1. Poll thinking chunk
        thinkPayload, _ := json.Marshal(map[string]interface{}{
          "operator_uri":     "@harpreetsahota/qwen38-max/qwen38_chat",
          "dataset_name":     datasetName,
          "view":             []interface{}{},
          "selected":         []interface{}{},
          "selected_samples": []interface{}{},
          "selected_labels":  []interface{}{},
          "filters":          map[string]interface{}{},
          "extended":         map[string]interface{}{},
          "params": map[string]interface{}{
            "panel_id":   "qwen38_chat",
            "__method__": "get_thinking_chunk",
            "run_id":     runID,
            "cursor":     thinkingCursor,
          },
        })
        tReq, err := http.NewRequestWithContext(r.Context(), "POST", executeURL, bytes.NewReader(thinkPayload))
        if err == nil {
          tReq.Header.Set("Content-Type", "application/json")
          tReq.Header.Set("X-FiftyOne-Session", token)
          tReq.Header.Set("User-Agent", reqUA)
          if reqHF != "" {
            tReq.Header.Set("Authorization", "Bearer "+reqHF)
          }
          tResp, tErr := gw.client.Do(tReq)
          if tErr == nil {
            var tResult struct {
              Result struct {
                Text   string `json:"text"`
                Cursor int    `json:"cursor"`
                Done   bool   `json:"done"`
              } `json:"result"`
            }
            if json.NewDecoder(tResp.Body).Decode(&tResult) == nil {
              if tResult.Result.Text != "" {
                streamer.Reasoning(tResult.Result.Text)
                thinkingCursor = tResult.Result.Cursor
              }
            }
            tResp.Body.Close()
          }
        }

        // 2. Poll stream chunk
        streamPayload, _ := json.Marshal(map[string]interface{}{
          "operator_uri":     "@harpreetsahota/qwen38-max/qwen38_chat",
          "dataset_name":     datasetName,
          "view":             []interface{}{},
          "selected":         []interface{}{},
          "selected_samples": []interface{}{},
          "selected_labels":  []interface{}{},
          "filters":          map[string]interface{}{},
          "extended":         map[string]interface{}{},
          "params": map[string]interface{}{
            "panel_id":   "qwen38_chat",
            "__method__": "get_stream_chunk",
            "run_id":     runID,
            "cursor":     streamCursor,
          },
        })
        sReq, err := http.NewRequestWithContext(r.Context(), "POST", executeURL, bytes.NewReader(streamPayload))
        if err == nil {
          sReq.Header.Set("Content-Type", "application/json")
          sReq.Header.Set("X-FiftyOne-Session", token)
          sReq.Header.Set("User-Agent", reqUA)
          if reqHF != "" {
            sReq.Header.Set("Authorization", "Bearer "+reqHF)
          }
          sResp, sErr := gw.client.Do(sReq)
          if sErr == nil {
            var sResult struct {
              Result struct {
                Text        string `json:"text"`
                Cursor      int    `json:"cursor"`
                Done        bool   `json:"done"`
                FinalStatus struct {
                  Status              string `json:"status"`
                  PromptTokens        int    `json:"prompt_tokens"`
                  CompletionTokens    int    `json:"completion_tokens"`
                  CachedTokens        int    `json:"cached_tokens"`
                  PromptTokensDetails struct {
                    CachedTokens int `json:"cached_tokens"`
                  } `json:"prompt_tokens_details"`
                  Error               string `json:"error"`
                } `json:"final_status"`
              } `json:"result"`
            }
            if json.NewDecoder(sResp.Body).Decode(&sResult) == nil {
              if sResult.Result.Text != "" {
                toolFilter.Feed(sResult.Result.Text)
                streamCursor = sResult.Result.Cursor
              }
              if sResult.Result.Done {
                sResp.Body.Close()
                toolFilter.Flush()

                finishReason := "stop"
                if toolFilter.HasEmittedTools() {
                  finishReason = "tool_calls"
                }
                streamer.Finish(finishReason)

                promptTokens := sResult.Result.FinalStatus.PromptTokens
                completionTokens := sResult.Result.FinalStatus.CompletionTokens
                cachedTokens := sResult.Result.FinalStatus.CachedTokens
                if cachedTokens == 0 {
                  cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens
                }
                cachedTokens = ResolveCachedTokens(cachedTokens, promptTokens, question)

                if req.StreamOptions != nil && req.StreamOptions.IncludeUsage {
                  streamer.EmitUsage(BuildUsage(promptTokens, completionTokens, cachedTokens, 0))
                }
                streamer.Done()
                log.Printf("[Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d cached=%d", runID, finishReason, promptTokens, completionTokens, cachedTokens)
                return
              }
            }
            sResp.Body.Close()
          }
        }
      }
    }
  } else {
    // Non-streaming mode: poll until done
    var accumulatedContent strings.Builder
    var accumulatedThinking strings.Builder
    promptTokens := 0
    completionTokens := 0

    streamCursor := 0
    thinkingCursor := 0
    pollTicker := time.NewTicker(200 * time.Millisecond)
    defer pollTicker.Stop()

    for {
      select {
      case <-r.Context().Done():
        log.Printf("[Non-Stream] Client canceled run_id=%s", runID)
        return
      case <-pollTicker.C:
        // Poll thinking chunk
        thinkPayload, _ := json.Marshal(map[string]interface{}{
          "operator_uri":     "@harpreetsahota/qwen38-max/qwen38_chat",
          "dataset_name":     datasetName,
          "view":             []interface{}{},
          "selected":         []interface{}{},
          "selected_samples": []interface{}{},
          "selected_labels":  []interface{}{},
          "filters":          map[string]interface{}{},
          "extended":         map[string]interface{}{},
          "params": map[string]interface{}{
            "panel_id":   "qwen38_chat",
            "__method__": "get_thinking_chunk",
            "run_id":     runID,
            "cursor":     thinkingCursor,
          },
        })
        tReq, err := http.NewRequestWithContext(r.Context(), "POST", executeURL, bytes.NewReader(thinkPayload))
        if err == nil {
          tReq.Header.Set("Content-Type", "application/json")
          tReq.Header.Set("X-FiftyOne-Session", token)
          tReq.Header.Set("User-Agent", reqUA)
          if reqHF != "" {
            tReq.Header.Set("Authorization", "Bearer "+reqHF)
          }
          tResp, tErr := gw.client.Do(tReq)
          if tErr == nil {
            var tResult struct {
              Result struct {
                Text   string `json:"text"`
                Cursor int    `json:"cursor"`
                Done   bool   `json:"done"`
              } `json:"result"`
            }
            if json.NewDecoder(tResp.Body).Decode(&tResult) == nil {
              if tResult.Result.Text != "" {
                accumulatedThinking.WriteString(tResult.Result.Text)
                thinkingCursor = tResult.Result.Cursor
              }
            }
            tResp.Body.Close()
          }
        }

        // Poll stream chunk
        streamPayload, _ := json.Marshal(map[string]interface{}{
          "operator_uri":     "@harpreetsahota/qwen38-max/qwen38_chat",
          "dataset_name":     datasetName,
          "view":             []interface{}{},
          "selected":         []interface{}{},
          "selected_samples": []interface{}{},
          "selected_labels":  []interface{}{},
          "filters":          map[string]interface{}{},
          "extended":         map[string]interface{}{},
          "params": map[string]interface{}{
            "panel_id":   "qwen38_chat",
            "__method__": "get_stream_chunk",
            "run_id":     runID,
            "cursor":     streamCursor,
          },
        })
        sReq, err := http.NewRequestWithContext(r.Context(), "POST", executeURL, bytes.NewReader(streamPayload))
        if err == nil {
          sReq.Header.Set("Content-Type", "application/json")
          sReq.Header.Set("X-FiftyOne-Session", token)
          sReq.Header.Set("User-Agent", reqUA)
          if reqHF != "" {
            sReq.Header.Set("Authorization", "Bearer "+reqHF)
          }
          sResp, sErr := gw.client.Do(sReq)
          if sErr == nil {
            var sResult struct {
              Result struct {
                Text        string `json:"text"`
                Cursor      int    `json:"cursor"`
                Done        bool   `json:"done"`
                FinalStatus struct {
                  Status              string `json:"status"`
                  PromptTokens        int    `json:"prompt_tokens"`
                  CompletionTokens    int    `json:"completion_tokens"`
                  CachedTokens        int    `json:"cached_tokens"`
                  PromptTokensDetails struct {
                    CachedTokens int `json:"cached_tokens"`
                  } `json:"prompt_tokens_details"`
                  Error               string `json:"error"`
                } `json:"final_status"`
              } `json:"result"`
            }
            if json.NewDecoder(sResp.Body).Decode(&sResult) == nil {
              if sResult.Result.Text != "" {
                accumulatedContent.WriteString(sResult.Result.Text)
                streamCursor = sResult.Result.Cursor
              }
              if sResult.Result.Done {
                sResp.Body.Close()
                promptTokens = sResult.Result.FinalStatus.PromptTokens
                completionTokens = sResult.Result.FinalStatus.CompletionTokens
                cachedTokens := sResult.Result.FinalStatus.CachedTokens
                if cachedTokens == 0 {
                  cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens
                }
                cachedTokens = ResolveCachedTokens(cachedTokens, promptTokens, question)

                fullContent := accumulatedContent.String()
                rawThinking := accumulatedThinking.String()

                // Check for embedded <think> tags in content
                embeddedThink, cleanedContent := ExtractThinkingContent(fullContent)
                if embeddedThink != "" {
                  if rawThinking != "" {
                    rawThinking = rawThinking + "\n\n" + embeddedThink
                  } else {
                    rawThinking = embeddedThink
                  }
                  fullContent = cleanedContent
                }

                toolCalls, remainingText := DetectToolCalls(fullContent)
                finishReason := "stop"
                var finalMsgContent interface{} = strings.TrimSpace(remainingText)

                if len(toolCalls) > 0 {
                  finishReason = "tool_calls"
                  if strings.TrimSpace(remainingText) == "" {
                    finalMsgContent = nil
                  }
                }

                out := FinalOutput{
                  Content:          finalMsgContent,
                  ReasoningContent: strings.TrimSpace(rawThinking),
                  ToolCalls:        toolCalls,
                  FinishReason:     finishReason,
                  PromptTokens:     promptTokens,
                  CompletionTokens: completionTokens,
                  CachedTokens:     cachedTokens,
                  ReasoningTokens:  0,
                }

                WriteCompletionResponse(w, completionID, createdTime, modelName, out)
                log.Printf("[Non-Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d cached=%d", runID, finishReason, promptTokens, completionTokens, cachedTokens)
                return
              }
            }
            sResp.Body.Close()
          }
        }
      }
    }
  }
}

// ---------------------------------------------------------------------------
// Gateway Entrypoint
// ---------------------------------------------------------------------------

func main() {
  port := flag.Int("port", 8080, "Port to listen on")
  spaceURL := flag.String("space-url", DefaultSpaceURL, "Upstream Hugging Face space URL")
  samplePath := flag.String("sample-path", DefaultSamplePath, "Target image sample filepath inside the container")
  modelID := flag.String("model", DefaultModelID, "Default model identifier")
  timeoutSec := flag.Int("timeout", 300, "Request timeout in seconds")
  userAge := flag.String("user-agent", "", "Custom User-Agent header")
  uaShort := flag.String("ua", "", "Custom User-Agent header (short)")
  hfToken := flag.String("hf-token", "", "Optional Hugging Face access token")
  flag.Parse()

  if *userAge != "" {
    ConfiguredUserAgent = *userAge
  } else if *uaShort != "" {
    ConfiguredUserAgent = *uaShort
  }
  if *hfToken != "" {
    ConfiguredToken = *hfToken
  }

  gw := NewQoronaGateway(*spaceURL, *samplePath, *modelID, time.Duration(*timeoutSec)*time.Second)

  mux := http.NewServeMux()

  // OpenAI routes
  mux.HandleFunc("/v1/models", gw.HandleModels)
  mux.HandleFunc("/models", gw.HandleModels)
  mux.HandleFunc("/v1/chat/completions", gw.HandleChatCompletions)
  mux.HandleFunc("/chat/completions", gw.HandleChatCompletions)

  // Health and index routes
  mux.HandleFunc("/__health", func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(`{"status":"ok","gateway":"qorona"}`))
  })
  mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(`{"status":"ok","gateway":"qorona"}`))
  })
  mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    EnableCORS(w)
    if r.URL.Path == "/" {
      w.Header().Set("Content-Type", "application/json")
      w.Write([]byte(`{"service":"qorona","description":"OpenAI-compatible gateway for Qwen 3.8 Max","models_endpoint":"/v1/models","completions_endpoint":"/v1/chat/completions"}`))
      return
    }
    http.NotFound(w, r)
  })

  addr := fmt.Sprintf(":%d", *port)
  log.Printf("[qorona] Listening on http://localhost%s (Upstream: %s)", addr, *spaceURL)
  log.Printf("[qorona] OpenAI compatible endpoints: http://localhost%s/v1/chat/completions", addr)

  if err := http.ListenAndServe(addr, mux); err != nil {
    log.Fatalf("[qorona] Server failed: %v", err)
  }
}