// groqqer: Pure Go OpenAI-compatible LLM gateway for the Groq Streamlit space
// Reverse engineers https://dromerosm-groq-chatbot.hf.space via direct WebSocket Protobuf wire format
// Created by Luxferre in 2026, released into the public domain with no warranties

package main

import (
  "bufio"
  "bytes"
  "context"
  "crypto/rand"
  "crypto/tls"
  "encoding/base64"
  "encoding/binary"
  "encoding/json"
  "errors"
  "flag"
  "fmt"
  "io"
  "log"
  "net"
  "net/http"
  "net/url"
  "os"
  "os/signal"
  "regexp"
  "strconv"
  "strings"
  "sync"
  "syscall"
  "time"
)

var (
  DefaultUserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
  DefaultTargetURL = "https://dromerosm-groq-chatbot.hf.space"
  DefaultModel     = "qwen/qwen3.6-27b"
)

// ---------------------------------------------------------------------------
// 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"`
  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"`
}

// ---------------------------------------------------------------------------
// UUID & Utility Functions
// ---------------------------------------------------------------------------

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

// ---------------------------------------------------------------------------
// SOCKS5 Proxy Client (RFC 1928 / RFC 1929)
// ---------------------------------------------------------------------------

type SOCKS5Config struct {
  Address  string
  Username string
  Password string
}

func ParseSOCKS5URL(proxyURL string) (*SOCKS5Config, error) {
  cleanURL := strings.TrimSpace(proxyURL)
  if cleanURL == "" {
    return nil, nil
  }
  if strings.HasPrefix(cleanURL, "socks5://") {
    cleanURL = strings.TrimPrefix(cleanURL, "socks5://")
  } else if strings.HasPrefix(cleanURL, "socks5h://") {
    cleanURL = strings.TrimPrefix(cleanURL, "socks5h://")
  }

  cfg := &SOCKS5Config{}
  if atIdx := strings.LastIndex(cleanURL, "@"); atIdx != -1 {
    userPass := cleanURL[:atIdx]
    cfg.Address = cleanURL[atIdx+1:]
    if colonIdx := strings.Index(userPass, ":"); colonIdx != -1 {
      cfg.Username = userPass[:colonIdx]
      cfg.Password = userPass[colonIdx+1:]
    } else {
      cfg.Username = userPass
    }
  } else {
    cfg.Address = cleanURL
  }

  if !strings.Contains(cfg.Address, ":") {
    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 {
    return nil, fmt.Errorf("invalid socks5 proxy configuration: %w", err)
  }
  if cfg == nil {
    var d net.Dialer
    return d.DialContext(ctx, "tcp", targetAddr)
  }

  var d net.Dialer
  conn, err := d.DialContext(ctx, "tcp", cfg.Address)
  if err != nil {
    return nil, fmt.Errorf("failed to connect to socks5 proxy at %s: %w", cfg.Address, err)
  }

  deadline, ok := ctx.Deadline()
  if !ok {
    deadline = time.Now().Add(30 * time.Second)
  }
  conn.SetDeadline(deadline)
  defer conn.SetDeadline(time.Time{})

  var greeting []byte
  if cfg.Username != "" {
    greeting = []byte{0x05, 0x02, 0x00, 0x02}
  } else {
    greeting = []byte{0x05, 0x01, 0x00}
  }

  if _, err := conn.Write(greeting); err != nil {
    conn.Close()
    return nil, fmt.Errorf("failed to write socks5 greeting: %w", err)
  }

  resp := make([]byte, 2)
  if _, err := io.ReadFull(conn, resp); err != nil {
    conn.Close()
    return nil, fmt.Errorf("failed to read socks5 greeting response: %w", err)
  }

  if resp[0] != 0x05 {
    conn.Close()
    return nil, fmt.Errorf("unsupported socks version: 0x%02x", resp[0])
  }

  if resp[1] == 0x02 {
    if cfg.Username == "" {
      conn.Close()
      return nil, fmt.Errorf("socks5 proxy requires authentication, but no credentials provided")
    }
    uLen := byte(len(cfg.Username))
    pLen := byte(len(cfg.Password))
    authReq := []byte{0x01, uLen}
    authReq = append(authReq, []byte(cfg.Username)...)
    authReq = append(authReq, pLen)
    authReq = append(authReq, []byte(cfg.Password)...)

    if _, err := conn.Write(authReq); err != nil {
      conn.Close()
      return nil, fmt.Errorf("failed to send socks5 authentication: %w", err)
    }

    authResp := make([]byte, 2)
    if _, err := io.ReadFull(conn, authResp); err != nil {
      conn.Close()
      return nil, fmt.Errorf("failed to read socks5 auth response: %w", err)
    }
    if authResp[1] != 0x00 {
      conn.Close()
      return nil, fmt.Errorf("socks5 authentication failed with status 0x%02x", authResp[1])
    }
  } else if resp[1] != 0x00 {
    conn.Close()
    return nil, fmt.Errorf("socks5 proxy rejected authentication methods: 0x%02x", resp[1])
  }

  host, portStr, err := net.SplitHostPort(targetAddr)
  if err != nil {
    conn.Close()
    return nil, fmt.Errorf("invalid target address %s: %w", targetAddr, err)
  }

  port, err := strconv.Atoi(portStr)
  if err != nil || port < 1 || port > 65535 {
    conn.Close()
    return nil, fmt.Errorf("invalid port in target address %s", targetAddr)
  }

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

  if _, err := conn.Write(reqBuf); err != nil {
    conn.Close()
    return nil, fmt.Errorf("failed to send socks5 connect request: %w", err)
  }

  respHdr := make([]byte, 4)
  if _, err := io.ReadFull(conn, respHdr); err != nil {
    conn.Close()
    return nil, fmt.Errorf("failed to read socks5 connect response: %w", err)
  }

  if respHdr[1] != 0x00 {
    conn.Close()
    return nil, fmt.Errorf("socks5 connect failed with reply code 0x%02x", respHdr[1])
  }

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

  return conn, nil
}

// ---------------------------------------------------------------------------
// Pure Go Protobuf Wire Formatter & Decoder
// ---------------------------------------------------------------------------

func encodeVarint(val uint64) []byte {
  var buf []byte
  for val >= 0x80 {
    buf = append(buf, byte(val|0x80))
    val >>= 7
  }
  buf = append(buf, byte(val))
  return buf
}

func encodeTag(fieldNum int, wireType int) []byte {
  return encodeVarint(uint64((fieldNum << 3) | wireType))
}

func encodeLengthDelimited(fieldNum int, data []byte) []byte {
  tag := encodeTag(fieldNum, 2)
  length := encodeVarint(uint64(len(data)))
  res := append(tag, length...)
  return append(res, data...)
}

func encodeString(fieldNum int, str string) []byte {
  return encodeLengthDelimited(fieldNum, []byte(str))
}

type ProtoField struct {
  Tag      int
  WireType int
  Varint   uint64
  Data     []byte
}

func decodeProtoFields(data []byte) []ProtoField {
  var fields []ProtoField
  r := bytes.NewReader(data)
  for r.Len() > 0 {
    rawTag, err := binary.ReadUvarint(r)
    if err != nil {
      break
    }
    fieldNum := int(rawTag >> 3)
    wireType := int(rawTag & 0x07)
    pf := ProtoField{Tag: fieldNum, WireType: wireType}
    switch wireType {
    case 0:
      v, err := binary.ReadUvarint(r)
      if err != nil {
        return fields
      }
      pf.Varint = v
    case 1:
      buf := make([]byte, 8)
      if _, err := io.ReadFull(r, buf); err != nil {
        return fields
      }
      pf.Data = buf
    case 2:
      length, err := binary.ReadUvarint(r)
      if err != nil {
        return fields
      }
      buf := make([]byte, length)
      if _, err := io.ReadFull(r, buf); err != nil {
        return fields
      }
      pf.Data = buf
    case 5:
      buf := make([]byte, 4)
      if _, err := io.ReadFull(r, buf); err != nil {
        return fields
      }
      pf.Data = buf
    default:
      return fields
    }
    fields = append(fields, pf)
  }
  return fields
}

type ParsedForwardMsg struct {
  IsFinished  bool
  DeltaPath   []uint64
  Markdown    string
  ChatInputID string
  SelectboxID string
  Options     []string
  AlertText   string
}

func parseForwardMsg(data []byte) ParsedForwardMsg {
  var res ParsedForwardMsg
  fields := decodeProtoFields(data)
  for _, f := range fields {
    if f.Tag == 6 && (f.Varint == 0 || f.Varint == 1 || f.Varint == 3 || f.Varint == 4 || f.Varint == 5) {
      res.IsFinished = true
    }
    if f.Tag == 2 && f.WireType == 2 { // ForwardMsgMetadata
      for _, mf := range decodeProtoFields(f.Data) {
        if mf.Tag == 2 { // delta_path
          if mf.WireType == 2 { // packed uint32
            r := bytes.NewReader(mf.Data)
            for r.Len() > 0 {
              v, _ := binary.ReadUvarint(r)
              res.DeltaPath = append(res.DeltaPath, v)
            }
          } else if mf.WireType == 0 {
            res.DeltaPath = append(res.DeltaPath, mf.Varint)
          }
        }
      }
    }
    if f.Tag == 5 && f.WireType == 2 { // Delta
      dFields := decodeProtoFields(f.Data)
      for _, df := range dFields {
        if df.Tag == 3 && df.WireType == 2 { // Element new_element
          eFields := decodeProtoFields(df.Data)
          for _, ef := range eFields {
            if ef.Tag == 29 && ef.WireType == 2 { // Markdown
              mFields := decodeProtoFields(ef.Data)
              for _, mf := range mFields {
                if mf.Tag == 1 && mf.WireType == 2 {
                  res.Markdown = string(mf.Data)
                }
              }
            } else if ef.Tag == 49 && ef.WireType == 2 { // ChatInput
              cFields := decodeProtoFields(ef.Data)
              for _, cf := range cFields {
                if cf.Tag == 1 && cf.WireType == 2 {
                  res.ChatInputID = string(cf.Data)
                }
              }
            } else if ef.Tag == 25 && ef.WireType == 2 { // Selectbox
              sFields := decodeProtoFields(ef.Data)
              for _, sf := range sFields {
                if sf.Tag == 1 && sf.WireType == 2 {
                  res.SelectboxID = string(sf.Data)
                } else if sf.Tag == 4 && sf.WireType == 2 {
                  res.Options = append(res.Options, string(sf.Data))
                }
              }
            } else if ef.Tag == 30 && ef.WireType == 2 { // Alert
              aFields := decodeProtoFields(ef.Data)
              for _, af := range aFields {
                if af.Tag == 1 && af.WireType == 2 {
                  res.AlertText = string(af.Data)
                }
              }
            }
          }
        }
      }
    }
  }
  return res
}

// ---------------------------------------------------------------------------
// RFC 6455 Pure Go WebSocket Client
// ---------------------------------------------------------------------------

func sendWSBinaryFrame(conn io.Writer, payload []byte) error {
  var header []byte
  header = append(header, 0x82) // Binary frame (0x02) | FIN (0x80)

  length := len(payload)
  maskKey := make([]byte, 4)
  _, _ = rand.Read(maskKey)

  if length < 126 {
    header = append(header, byte(length|0x80))
  } else if length < 65536 {
    header = append(header, 126|0x80)
    var b [2]byte
    binary.BigEndian.PutUint16(b[:], uint16(length))
    header = append(header, b[:]...)
  } else {
    header = append(header, 127|0x80)
    var b [8]byte
    binary.BigEndian.PutUint64(b[:], uint64(length))
    header = append(header, b[:]...)
  }

  header = append(header, maskKey...)
  masked := make([]byte, length)
  for i := 0; i < length; i++ {
    masked[i] = payload[i] ^ maskKey[i%4]
  }

  _, err := conn.Write(append(header, masked...))
  return err
}

func sendWSPongFrame(conn io.Writer, payload []byte) error {
  var header []byte
  header = append(header, 0x8a) // Pong (0x0A) | FIN (0x80)

  length := len(payload)
  maskKey := make([]byte, 4)
  _, _ = rand.Read(maskKey)

  header = append(header, byte(length|0x80))
  header = append(header, maskKey...)
  masked := make([]byte, length)
  for i := 0; i < length; i++ {
    masked[i] = payload[i] ^ maskKey[i%4]
  }

  _, err := conn.Write(append(header, masked...))
  return err
}

func readWSFrame(conn net.Conn, reader *bufio.Reader, timeout time.Duration) ([]byte, byte, error) {
  if timeout > 0 {
    _ = conn.SetReadDeadline(time.Now().Add(timeout))
  } else {
    _ = conn.SetReadDeadline(time.Time{})
  }

  b1, err := reader.ReadByte()
  if err != nil {
    return nil, 0, err
  }
  b2, err := reader.ReadByte()
  if err != nil {
    return nil, 0, err
  }

  opcode := b1 & 0x0f
  isMasked := (b2 & 0x80) != 0
  length := int(b2 & 0x7f)

  if length == 126 {
    var ext uint16
    if err := binary.Read(reader, binary.BigEndian, &ext); err != nil {
      return nil, 0, err
    }
    length = int(ext)
  } else if length == 127 {
    var ext uint64
    if err := binary.Read(reader, binary.BigEndian, &ext); err != nil {
      return nil, 0, err
    }
    length = int(ext)
  }

  var maskKey [4]byte
  if isMasked {
    if _, err := io.ReadFull(reader, maskKey[:]); err != nil {
      return nil, 0, err
    }
  }

  payload := make([]byte, length)
  if _, err := io.ReadFull(reader, payload); err != nil {
    return nil, 0, err
  }

  if isMasked {
    for i := 0; i < length; i++ {
      payload[i] ^= maskKey[i%4]
    }
  }

  return payload, opcode, nil
}

// ---------------------------------------------------------------------------
// Streamlit Session & Client
// ---------------------------------------------------------------------------

type StreamlitSession struct {
  conn        net.Conn
  reader      *bufio.Reader
  chatInputID string
  selectboxID string
  options     []string
  activeModel string
  closed      bool
  mu          sync.Mutex
}

func (s *StreamlitSession) Close() {
  s.mu.Lock()
  defer s.mu.Unlock()
  if !s.closed {
    s.closed = true
    if s.conn != nil {
      _ = s.conn.Close()
    }
  }
}

func dialStreamlitWebSocket(ctx context.Context, targetURL, proxyURL, userAgent string) (net.Conn, *bufio.Reader, error) {
  u, err := url.Parse(targetURL)
  if err != nil {
    return nil, nil, fmt.Errorf("invalid target url: %w", err)
  }

  host := u.Hostname()
  port := u.Port()
  isTLS := u.Scheme == "https" || u.Scheme == "wss"
  if port == "" {
    if isTLS {
      port = "443"
    } else {
      port = "80"
    }
  }

  targetAddr := net.JoinHostPort(host, port)
  var rawConn net.Conn

  if proxyURL != "" {
    rawConn, err = DialSOCKS5(ctx, proxyURL, targetAddr)
  } else {
    var d net.Dialer
    rawConn, err = d.DialContext(ctx, "tcp", targetAddr)
  }
  if err != nil {
    return nil, nil, fmt.Errorf("failed to connect to %s: %w", targetAddr, err)
  }

  var conn net.Conn = rawConn
  if isTLS {
    tlsConfig := &tls.Config{
      ServerName: host,
    }
    tlsConn := tls.Client(rawConn, tlsConfig)
    if err := tlsConn.HandshakeContext(ctx); err != nil {
      rawConn.Close()
      return nil, nil, fmt.Errorf("tls handshake failed: %w", err)
    }
    conn = tlsConn
  }

  keyBytes := make([]byte, 16)
  _, _ = rand.Read(keyBytes)
  secKey := base64.StdEncoding.EncodeToString(keyBytes)

  originScheme := "https"
  if !isTLS {
    originScheme = "http"
  }
  origin := fmt.Sprintf("%s://%s", originScheme, host)

  req := fmt.Sprintf(
    "GET /_stcore/stream HTTP/1.1\r\n"+
      "Host: %s\r\n"+
      "Upgrade: websocket\r\n"+
      "Connection: Upgrade\r\n"+
      "Sec-WebSocket-Key: %s\r\n"+
      "Sec-WebSocket-Version: 13\r\n"+
      "Origin: %s\r\n"+
      "User-Agent: %s\r\n\r\n",
    host, secKey, origin, userAgent,
  )

  if _, err := conn.Write([]byte(req)); err != nil {
    conn.Close()
    return nil, nil, fmt.Errorf("failed to write websocket handshake request: %w", err)
  }

  reader := bufio.NewReader(conn)
  statusLine, err := reader.ReadString('\n')
  if err != nil || !strings.Contains(statusLine, "101") {
    conn.Close()
    return nil, nil, fmt.Errorf("websocket handshake rejected: %s", statusLine)
  }

  for {
    line, err := reader.ReadString('\n')
    if err != nil || strings.TrimSpace(line) == "" {
      break
    }
  }

  return conn, reader, nil
}

func OpenStreamlitSession(ctx context.Context, targetURL, proxyURL, userAgent string) (*StreamlitSession, error) {
  conn, reader, err := dialStreamlitWebSocket(ctx, targetURL, proxyURL, userAgent)
  if err != nil {
    return nil, err
  }

  sess := &StreamlitSession{
    conn:   conn,
    reader: reader,
  }

  // Initial rerun_script BackMsg (ClientState: query_string = "")
  clientState := encodeString(1, "")
  backMsg := encodeLengthDelimited(11, clientState)

  if err := sendWSBinaryFrame(conn, backMsg); err != nil {
    sess.Close()
    return nil, fmt.Errorf("failed to send initial rerun BackMsg: %w", err)
  }

  // Read frames until script_finished
  for {
    payload, opcode, err := readWSFrame(conn, reader, 25*time.Second)
    if err != nil {
      sess.Close()
      return nil, fmt.Errorf("error reading initial session frames: %w", err)
    }

    if opcode == 0x09 { // Ping
      _ = sendWSPongFrame(conn, payload)
      continue
    }

    parsed := parseForwardMsg(payload)
    if parsed.ChatInputID != "" {
      sess.chatInputID = parsed.ChatInputID
    }
    if parsed.SelectboxID != "" {
      sess.selectboxID = parsed.SelectboxID
      sess.options = parsed.Options
    }
    if parsed.IsFinished {
      break
    }
  }

  if sess.chatInputID == "" {
    // Default fallback chatInput widget ID if not matched in stream
    sess.chatInputID = "$$ID-ae617304c8297cf4ddb1a23ee392ce5b-None"
  }
  if sess.selectboxID == "" {
    sess.selectboxID = "$$ID-625687b40312ff8d73403e7858a2f552-chat_model_selector"
  }

  return sess, nil
}

func (s *StreamlitSession) SwitchModel(modelID string) error {
  if modelID == "" || s.selectboxID == "" {
    return nil
  }

  cleanModel := strings.TrimSpace(modelID)
  if idx := strings.Index(cleanModel, " ("); idx != -1 {
    cleanModel = strings.TrimSpace(cleanModel[:idx])
  }

  // Match model against options
  matchedModel := cleanModel
  for _, opt := range s.options {
    rawOpt := opt
    if idx := strings.Index(rawOpt, " ("); idx != -1 {
      rawOpt = strings.TrimSpace(rawOpt[:idx])
    }
    if strings.EqualFold(rawOpt, cleanModel) || strings.Contains(strings.ToLower(rawOpt), strings.ToLower(cleanModel)) {
      matchedModel = rawOpt
      break
    }
  }

  if s.activeModel == matchedModel {
    return nil
  }

  // Streamlit Selectbox sends string_value (field 6) in WidgetState
  sbWidget := append(encodeString(1, s.selectboxID), encodeString(6, matchedModel)...)
  wsData := encodeLengthDelimited(1, sbWidget)
  wStates := encodeLengthDelimited(2, wsData)
  csData := append(encodeString(1, ""), wStates...)
  switchBackMsg := encodeLengthDelimited(11, csData)

  if err := sendWSBinaryFrame(s.conn, switchBackMsg); err != nil {
    return fmt.Errorf("failed to send model switch BackMsg: %w", err)
  }

  for {
    payload, opcode, err := readWSFrame(s.conn, s.reader, 25*time.Second)
    if err != nil {
      return fmt.Errorf("error reading model switch frames: %w", err)
    }
    if opcode == 0x09 {
      _ = sendWSPongFrame(s.conn, payload)
      continue
    }
    parsed := parseForwardMsg(payload)
    if parsed.ChatInputID != "" {
      s.chatInputID = parsed.ChatInputID
    }
    if parsed.SelectboxID != "" {
      s.selectboxID = parsed.SelectboxID
    }
    if parsed.IsFinished {
      break
    }
  }

  s.activeModel = matchedModel
  return nil
}

type GatewayAlertError struct {
  Alert       string
  IsRateLimit bool
  RetryAfter  int
}

func (e *GatewayAlertError) Error() string {
  return e.Alert
}

func parseRetryAfter(alert string) int {
  re := regexp.MustCompile(`(?i)(?:retry\s+after|retry\s+in)\s+(\d+)`)
  match := re.FindStringSubmatch(alert)
  if len(match) >= 2 {
    sec, _ := strconv.Atoi(match[1])
    if sec > 0 {
      return sec
    }
  }
  return 6
}

func isRateLimit(alert string) bool {
  lower := strings.ToLower(alert)
  return strings.Contains(lower, "rate limit") || strings.Contains(lower, "too many requests") || strings.Contains(lower, "429")
}

func resolveAlert(alerts []string) (alertText string, isRL bool, retryAfter int) {
  for _, a := range alerts {
    if isRateLimit(a) {
      return a, true, parseRetryAfter(a)
    }
  }

  for _, a := range alerts {
    lower := strings.ToLower(a)
    if !strings.Contains(lower, "the last response did not finish") {
      return a, false, 0
    }
  }

  if len(alerts) > 0 {
    return alerts[len(alerts)-1], false, 0
  }
  return "", false, 0
}

func isStaticUIMarkdown(text, promptText string) bool {
  trimmed := strings.TrimSpace(text)
  if trimmed == "" || trimmed == strings.TrimSpace(promptText) {
    return true
  }
  staticPrefixes := []string{
    "Model:",
    "Model :",
    "- **Model ID:**",
    "[Powered by Groq]",
    "---",
    "-----",
    "Text models available",
    "This model accepts and returns text",
    "This model uses non-streaming responses",
    "Select a vision-capable model",
    "Attached document:",
    "### Upload an Image",
    "### Usage Summary",
    "### Chat Interface",
    "### Upload Context",
    "### Comparison Context",
    "A document is already attached",
    "**Token Usage",
    "**Latest API Rate Snapshot:**",
    "**Request Max Tokens:**",
    "**Conversation Total:**",
    "**Important:**",
    "Important:",
    "Partial response: generation did not finish.",
  }
  for _, p := range staticPrefixes {
    if strings.HasPrefix(trimmed, p) {
      return true
    }
  }
  if strings.Contains(trimmed, "**Important:**") ||
    strings.Contains(trimmed, "Important:") ||
    strings.Contains(trimmed, "older messages were left out of this request") ||
    strings.Contains(trimmed, "Response limit adjusted to") {
    return true
  }
  return false
}

func (s *StreamlitSession) SubmitPrompt(promptText string, onChunk func(delta string), onAlert func(alert string)) error {
  // Construct BackMsg with chat_input_value and selected model in selectbox
  chatInputVal := encodeString(1, promptText)
  ciWidget := append(encodeString(1, s.chatInputID), encodeLengthDelimited(15, chatInputVal)...)

  var wsData []byte
  if s.selectboxID != "" && s.activeModel != "" {
    sbWidget := append(encodeString(1, s.selectboxID), encodeString(6, s.activeModel)...)
    wsData = append(encodeLengthDelimited(1, sbWidget), encodeLengthDelimited(1, ciWidget)...)
  } else {
    wsData = encodeLengthDelimited(1, ciWidget)
  }

  wStates := encodeLengthDelimited(2, wsData)
  csData := append(encodeString(1, ""), wStates...)
  promptBackMsg := encodeLengthDelimited(11, csData)

  if err := sendWSBinaryFrame(s.conn, promptBackMsg); err != nil {
    return fmt.Errorf("failed to send prompt BackMsg: %w", err)
  }

  lastLen := 0
  assistantPathKey := ""
  userEchoSeen := false
  var alerts []string

  for {
    payload, opcode, err := readWSFrame(s.conn, s.reader, 90*time.Second)
    if err != nil {
      return fmt.Errorf("connection closed during response: %w", err)
    }

    if opcode == 0x09 { // Ping
      _ = sendWSPongFrame(s.conn, payload)
      continue
    }

    parsed := parseForwardMsg(payload)

    if parsed.AlertText != "" {
      alerts = append(alerts, parsed.AlertText)
      if onAlert != nil {
        onAlert(parsed.AlertText)
      }
    }

    if parsed.Markdown != "" {
      // Ignore sidebar elements (deltaPath[0] == 1)
      if len(parsed.DeltaPath) > 0 && parsed.DeltaPath[0] == 1 {
        continue
      }

      // In this Streamlit space:
      // DeltaPath[0] == 0 is the main page.
      // DeltaPath[1] <= 5 are the header widgets (0-4) and the user chat echo container (5).
      // The assistant response is always DeltaPath[1] >= 6.
      if len(parsed.DeltaPath) >= 2 && parsed.DeltaPath[0] == 0 && parsed.DeltaPath[1] <= 5 {
        continue
      }

      // Metadata or captions inside message containers have index >= 1 (e.g. [0, 6, 1])
      if len(parsed.DeltaPath) >= 3 && parsed.DeltaPath[len(parsed.DeltaPath)-1] != 0 {
        continue
      }

      // Check if user echo arrived (fallback for environments without hierarchical delta paths)
      trimmedMD := strings.TrimSpace(parsed.Markdown)
      trimmedPrompt := strings.TrimSpace(promptText)
      if !userEchoSeen {
        if trimmedMD == trimmedPrompt ||
          strings.HasPrefix(trimmedPrompt, trimmedMD) ||
          strings.HasPrefix(trimmedMD, trimmedPrompt) {
          userEchoSeen = true
          continue
        }
      }

      // If text is not static UI, stream assistant delta
      if !isStaticUIMarkdown(parsed.Markdown, promptText) {
        pathKey := fmt.Sprintf("%v", parsed.DeltaPath)
        if assistantPathKey == "" {
          assistantPathKey = pathKey
        }
        if pathKey == assistantPathKey {
          if len(parsed.Markdown) > lastLen {
            delta := parsed.Markdown[lastLen:]
            lastLen = len(parsed.Markdown)
            if onChunk != nil {
              onChunk(delta)
            }
          }
        }
      }
    }

    if parsed.IsFinished {
      break
    }
  }

  if lastLen == 0 {
    if alertText, isRL, retrySec := resolveAlert(alerts); alertText != "" {
      return &GatewayAlertError{
        Alert:       alertText,
        IsRateLimit: isRL,
        RetryAfter:  retrySec,
      }
    }
    return fmt.Errorf("empty assistant response received from space")
  }

  return nil
}

// ---------------------------------------------------------------------------
// Tool and Message Processing
// ---------------------------------------------------------------------------

func BuildToolInstruction(tools []Tool, toolChoice interface{}) string {
  if len(tools) == 0 {
    return ""
  }
  toolsBytes, _ := json.Marshal(tools)
  var sb strings.Builder
  sb.WriteString("\n\n# Tool Calling Instructions\nYou have access to the following functions:\n<tools>\n")
  sb.WriteString(string(toolsBytes))
  sb.WriteString("\n</tools>\n\n")
  sb.WriteString("When you need to call a function, respond ONLY with a <tool_call> block formatted exactly as follows:\n")
  sb.WriteString("<tool_call>\n{\"name\": \"<function-name>\", \"arguments\": {<args-json-object>}}\n</tool_call>\n\n")
  sb.WriteString("CRITICAL EXECUTION RULES:\n")
  sb.WriteString("1. If you invoke a tool, output ONLY the <tool_call> block. Do not write conversational greetings, explanations, or filler text outside the tags.\n")
  sb.WriteString("2. Put any reasoning or thought process inside <think>...</think> tags.\n")
  sb.WriteString("3. If multiple tools need to be called, output each inside its own <tool_call>...</tool_call> block.\n")
  sb.WriteString("4. When tool execution results are provided in <tool_response> blocks, inspect the output:\n")
  sb.WriteString("   - If further steps or additional tools are needed to fulfill the user's request, emit the next <tool_call> block.\n")
  sb.WriteString("   - If all needed information has been retrieved, synthesize the answers and deliver the final response to the user.\n")
  sb.WriteString("   - Never stop or terminate the conversation prematurely while intermediate steps remain.\n")

  if toolChoice != nil {
    if choiceStr, ok := toolChoice.(string); ok {
      if choiceStr == "required" {
        sb.WriteString("\n\nYou MUST invoke at least one tool from the list above.")
      }
    } else if choiceMap, ok := toolChoice.(map[string]interface{}); ok {
      if fnMap, ok := choiceMap["function"].(map[string]interface{}); ok {
        if fnName, ok := fnMap["name"].(string); ok && fnName != "" {
          sb.WriteString(fmt.Sprintf("\n\nYou MUST call the function %q.", fnName))
        }
      }
    }
  }

  return sb.String()
}

func ExtractToolName(fn interface{}) string {
  if fn == nil {
    return ""
  }
  if fnMap, ok := fn.(map[string]interface{}); ok {
    if name, ok := fnMap["name"].(string); ok {
      return strings.TrimSpace(name)
    }
  }
  if tcf, ok := fn.(ToolCallFunction); ok {
    return strings.TrimSpace(tcf.Name)
  }
  b, err := json.Marshal(fn)
  if err == nil {
    var m map[string]interface{}
    if json.Unmarshal(b, &m) == nil {
      if name, ok := m["name"].(string); ok {
        return strings.TrimSpace(name)
      }
    }
  }
  return ""
}

func GetAllowedToolNames(tools []Tool, toolChoice interface{}) map[string]bool {
  if len(tools) == 0 {
    return nil
  }

  if choiceStr, ok := toolChoice.(string); ok {
    if choiceStr == "none" {
      return nil
    }
  }

  if choiceMap, ok := toolChoice.(map[string]interface{}); ok {
    if fnMap, ok := choiceMap["function"].(map[string]interface{}); ok {
      if name, ok := fnMap["name"].(string); ok && strings.TrimSpace(name) != "" {
        return map[string]bool{strings.TrimSpace(name): true}
      }
    }
  }

  allowed := make(map[string]bool)
  for _, t := range tools {
    name := ExtractToolName(t.Function)
    if name != "" {
      allowed[name] = true
    }
  }

  if len(allowed) == 0 {
    return nil
  }
  return allowed
}

var toolCallTagRegex = regexp.MustCompile(`(?s)<tool_call>.*?</tool_call>`)

func stripToolCallTags(s string) string {
  return strings.TrimSpace(toolCallTagRegex.ReplaceAllString(s, ""))
}

func FormatPrompt(req ChatCompletionRequest) string {
  toolInstruction := BuildToolInstruction(req.Tools, req.ToolChoice)

  // Build mapping from tool_call_id to function name
  toolCallIDToName := make(map[string]string)
  for _, msg := range req.Messages {
    if msg.Role == "assistant" {
      for _, tc := range msg.ToolCalls {
        if tc.ID != "" && tc.Function.Name != "" {
          toolCallIDToName[tc.ID] = tc.Function.Name
        }
      }
    }
  }

  // Find where trailing tool responses start (if any)
  firstTrailingToolIdx := len(req.Messages)
  for i := len(req.Messages) - 1; i >= 0; i-- {
    role := req.Messages[i].Role
    if role == "tool" || role == "function" {
      firstTrailingToolIdx = i
    } else {
      break
    }
  }

  var systemInstructions []string
  var historyTurns []string
  var currentTurn string

  // Process messages before the trailing tool responses
  endHistoryIdx := firstTrailingToolIdx
  if firstTrailingToolIdx == len(req.Messages) && len(req.Messages) > 0 {
    if req.Messages[len(req.Messages)-1].Role == "user" {
      endHistoryIdx = len(req.Messages) - 1
      currentTurn = req.Messages[len(req.Messages)-1].GetContentString()
    }
  }

  for i := 0; i < endHistoryIdx; i++ {
    msg := req.Messages[i]
    contentStr := msg.GetContentString()
    switch msg.Role {
    case "system":
      if contentStr != "" {
        systemInstructions = append(systemInstructions, contentStr)
      }
    case "assistant":
      cleanAssistant := strings.TrimSpace(contentStr)
      if strings.HasPrefix(cleanAssistant, "Model:") || strings.HasPrefix(cleanAssistant, "Model :") {
        cleanAssistant = ""
      }
      if len(msg.ToolCalls) > 0 {
        cleanAssistant = stripToolCallTags(cleanAssistant)
      }
      var sb strings.Builder
      if cleanAssistant != "" {
        sb.WriteString(cleanAssistant)
      }
      for _, tc := range msg.ToolCalls {
        if sb.Len() > 0 {
          sb.WriteString("\n")
        }
        args := tc.Function.Arguments
        if strings.TrimSpace(args) == "" {
          args = "{}"
        }
        sb.WriteString(fmt.Sprintf("<tool_call>\n{\"name\": %q, \"arguments\": %s}\n</tool_call>", tc.Function.Name, args))
      }
      if sb.Len() > 0 {
        historyTurns = append(historyTurns, "Assistant: "+sb.String())
      }
    case "tool", "function":
      toolName := msg.Name
      if resolved, ok := toolCallIDToName[msg.ToolCallID]; ok && resolved != "" {
        toolName = resolved
      } else if toolName == "" {
        toolName = msg.ToolCallID
      }
      var contentJSON []byte
      if json.Valid([]byte(contentStr)) {
        contentJSON = []byte(contentStr)
      } else {
        contentJSON, _ = json.Marshal(contentStr)
      }
      jsonStr := string(contentJSON)
      if len(jsonStr) > 2000 {
        jsonStr = jsonStr[:1800] + `"... [truncated]"`
      }
      var turnText string
      if msg.ToolCallID != "" {
        turnText = fmt.Sprintf("<tool_response>\n{\"name\": %q, \"tool_call_id\": %q, \"content\": %s}\n</tool_response>", toolName, msg.ToolCallID, jsonStr)
      } else {
        turnText = fmt.Sprintf("<tool_response>\n{\"name\": %q, \"content\": %s}\n</tool_response>", toolName, jsonStr)
      }
      historyTurns = append(historyTurns, fmt.Sprintf("Tool Result (%s): %s", toolName, turnText))
    case "user":
      historyTurns = append(historyTurns, "User: "+contentStr)
    }
  }

  // If we have trailing tool responses, group them together with a continuation directive
  if firstTrailingToolIdx < len(req.Messages) {
    var toolResSb strings.Builder
    toolResSb.WriteString("[Tool Execution Results]\n")
    for i := firstTrailingToolIdx; i < len(req.Messages); i++ {
      msg := req.Messages[i]
      toolName := msg.Name
      if resolved, ok := toolCallIDToName[msg.ToolCallID]; ok && resolved != "" {
        toolName = resolved
      } else if toolName == "" {
        toolName = msg.ToolCallID
      }
      contentStr := msg.GetContentString()
      var contentJSON []byte
      if json.Valid([]byte(contentStr)) {
        contentJSON = []byte(contentStr)
      } else {
        contentJSON, _ = json.Marshal(contentStr)
      }
      jsonStr := string(contentJSON)
      if len(jsonStr) > 3500 {
        jsonStr = jsonStr[:3200] + `"... [truncated]"`
      }

      toolResSb.WriteString("<tool_response>\n")
      if msg.ToolCallID != "" {
        toolResSb.WriteString(fmt.Sprintf("{\"name\": %q, \"tool_call_id\": %q, \"content\": %s}\n", toolName, msg.ToolCallID, jsonStr))
      } else {
        toolResSb.WriteString(fmt.Sprintf("{\"name\": %q, \"content\": %s}\n", toolName, jsonStr))
      }
      toolResSb.WriteString("</tool_response>\n")
    }
    toolResSb.WriteString("\n[Next Steps Directive]\n")
    toolResSb.WriteString("You have received the results of the tool execution(s) above.\n")
    toolResSb.WriteString("- Analyze these results in the context of the user request and conversation history.\n")
    toolResSb.WriteString("- If additional tool calls are needed to complete the task, emit the next <tool_call> block immediately.\n")
    toolResSb.WriteString("- If all information needed to fulfill the request is now available, provide the final comprehensive response to the user.\n")
    toolResSb.WriteString("- Do not terminate or stop without answering the user or taking the next step.")
    currentTurn = toolResSb.String()
  }

  // Prune history turns if history exceeds budget (~12000 chars)
  totalHistLen := 0
  for _, t := range historyTurns {
    totalHistLen += len(t)
  }
  if totalHistLen > 12000 && len(historyTurns) > 3 {
    firstTurn := historyTurns[0]
    var pruned []string
    pruned = append(pruned, firstTurn)
    pruned = append(pruned, "[... earlier conversation turns omitted for brevity ...]")
    budget := 10000
    var recent []string
    currentBudget := 0
    for i := len(historyTurns) - 1; i >= 1; i-- {
      tLen := len(historyTurns[i])
      if currentBudget+tLen > budget {
        break
      }
      currentBudget += tLen
      recent = append([]string{historyTurns[i]}, recent...)
    }
    pruned = append(pruned, recent...)
    historyTurns = pruned
  }

  var promptBuilder strings.Builder

  if len(systemInstructions) > 0 || toolInstruction != "" {
    promptBuilder.WriteString("[System Instructions]\n")
    for _, sys := range systemInstructions {
      promptBuilder.WriteString(sys)
      promptBuilder.WriteString("\n")
    }
    if toolInstruction != "" {
      promptBuilder.WriteString(strings.TrimSpace(toolInstruction))
      promptBuilder.WriteString("\n")
    }
    promptBuilder.WriteString("\n")
  }

  if len(historyTurns) > 0 {
    promptBuilder.WriteString("[Conversation History]\n")
    for _, turn := range historyTurns {
      promptBuilder.WriteString(turn)
      promptBuilder.WriteString("\n")
    }
    promptBuilder.WriteString("\n")
  }

  if currentTurn != "" {
    if strings.HasPrefix(currentTurn, "[Tool Execution Results]") {
      promptBuilder.WriteString(currentTurn)
    } else {
      if len(historyTurns) > 0 || len(systemInstructions) > 0 {
        promptBuilder.WriteString("[User]\n")
      }
      promptBuilder.WriteString(currentTurn)
    }
  }

  return strings.TrimSpace(promptBuilder.String())
}

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
  }
}

var toolNameRegex = regexp.MustCompile(`"\s*(?:name|function|action|call)\s*"\s*:\s*"\s*([^"]+?)\s*"`)

func repairToolCallJSON(jsonStr string) (ToolCall, bool) {
  nameMatch := toolNameRegex.FindStringSubmatch(jsonStr)
  if len(nameMatch) < 2 {
    return ToolCall{}, false
  }
  nameVal := strings.TrimSpace(nameMatch[1])

  argsStr := "{}"
  argsKwList := []string{`"arguments"`, `" parameters "`, `"arguments "`, `" parameters"`, `"parameters"`, `"args"`, `"input"`}
  argsIdx := -1
  for _, kw := range argsKwList {
    idx := strings.Index(jsonStr, kw)
    if idx >= 0 {
      argsIdx = idx + len(kw)
      break
    }
  }

  var targetStr string
  if argsIdx >= 0 {
    targetStr = strings.TrimSpace(jsonStr[argsIdx:])
    if strings.HasPrefix(targetStr, ":") {
      targetStr = strings.TrimSpace(targetStr[1:])
    }
  } else {
    targetStr = jsonStr
  }

  if strings.HasPrefix(targetStr, "{") {
    endIdx := strings.LastIndex(targetStr, "}")
    if endIdx > 0 {
      objCandidate := targetStr[:endIdx+1]
      var testMap map[string]interface{}
      if json.Unmarshal([]byte(objCandidate), &testMap) == nil {
        b, _ := json.Marshal(sanitizeJSONValue(testMap))
        return ToolCall{
          ID:   "call_" + GenerateUUID()[:8],
          Type: "function",
          Function: ToolCallFunction{
            Name:      nameVal,
            Arguments: string(b),
          },
        }, true
      }
    }
  } else if strings.HasPrefix(targetStr, `"`) {
    endIdx := strings.LastIndex(targetStr, `"`)
    if endIdx > 0 {
      val := strings.TrimSpace(targetStr[1:endIdx])
      b, _ := json.Marshal(val)
      argsStr = string(b)
    }
  }

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

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 repairToolCallJSON(cleaned)
}

func parseToolCallsFromText(input string) ([]ToolCall, bool) {
  cleaned := cleanJSONBlock(strings.TrimSpace(input))
  if cleaned == "" {
    return nil, false
  }

  // Check if it is a JSON array of tool calls: [...]
  if strings.HasPrefix(cleaned, "[") && strings.HasSuffix(cleaned, "]") {
    var rawSlice []interface{}
    if err := json.Unmarshal([]byte(cleaned), &rawSlice); err == nil && len(rawSlice) > 0 {
      var calls []ToolCall
      for _, item := range rawSlice {
        b, _ := json.Marshal(item)
        if tc, ok := parseSingleToolCall(string(b)); ok {
          calls = append(calls, tc)
        } else if tc, ok := repairToolCallJSON(string(b)); ok {
          calls = append(calls, tc)
        }
      }
      if len(calls) > 0 {
        return calls, true
      }
    }
  }

  // Check single tool call
  if tc, ok := parseSingleToolCall(cleaned); ok {
    return []ToolCall{tc}, true
  }

  return nil, 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 calls, ok := parseToolCallsFromText(inner); ok {
    return calls, 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])
    }
  }

  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 nil, false
}

func ExtractToolCallBlocks(content string) (blocks []string, remaining string) {
  var b []string
  s := content
  searchOffset := 0
  var remainingBuilder strings.Builder
  lastPos := 0

  for {
    relStart := strings.Index(s[searchOffset:], "<tool_call>")
    if relStart == -1 {
      break
    }
    startIdx := searchOffset + relStart
    afterStart := startIdx + len("<tool_call>")

    relEnd := strings.Index(s[afterStart:], "</tool_call>")
    if relEnd == -1 {
      // Unclosed <tool_call>, do not treat as block
      searchOffset = afterStart
      continue
    }

    if innerNext := strings.Index(s[afterStart:afterStart+relEnd], "<tool_call>"); innerNext != -1 {
      searchOffset = afterStart + innerNext
      continue
    }

    endIdx := afterStart + relEnd + len("</tool_call>")
    blockText := s[startIdx:endIdx]
    b = append(b, blockText)

    if startIdx > lastPos {
      remainingBuilder.WriteString(s[lastPos:startIdx])
    }
    lastPos = endIdx
    searchOffset = endIdx
  }

  if lastPos < len(s) {
    remainingBuilder.WriteString(s[lastPos:])
  }

  return b, strings.TrimSpace(remainingBuilder.String())
}

func DetectToolCalls(content string, allowedTools map[string]bool) ([]ToolCall, string, bool) {
  if len(allowedTools) == 0 {
    return nil, content, false
  }

  type validBlock struct {
    startIdx int
    endIdx   int
    calls    []ToolCall
  }

  var validBlocks []validBlock
  searchOffset := 0

  for {
    relStart := strings.Index(content[searchOffset:], "<tool_call>")
    if relStart == -1 {
      break
    }
    startIdx := searchOffset + relStart
    afterStart := startIdx + len("<tool_call>")

    relEnd := strings.Index(content[afterStart:], "</tool_call>")
    if relEnd == -1 {
      // Unclosed <tool_call> tag - leave as plain content
      searchOffset = afterStart
      continue
    }

    if innerNext := strings.Index(content[afterStart:afterStart+relEnd], "<tool_call>"); innerNext != -1 {
      searchOffset = afterStart + innerNext
      continue
    }

    endIdx := afterStart + relEnd + len("</tool_call>")
    blockText := content[startIdx:endIdx]

    if parsedCalls, ok := parseXMLToolCall(blockText); ok {
      var callsForBlock []ToolCall
      for _, c := range parsedCalls {
        if allowedTools[c.Function.Name] {
          callsForBlock = append(callsForBlock, c)
        }
      }
      if len(callsForBlock) > 0 {
        validBlocks = append(validBlocks, validBlock{
          startIdx: startIdx,
          endIdx:   endIdx,
          calls:    callsForBlock,
        })
        searchOffset = endIdx
        continue
      }
    }

    // Not a valid tool call or function name not allowed - keep in content
    searchOffset = endIdx
  }

  if len(validBlocks) > 0 {
    var allCalls []ToolCall
    var remainingBuilder strings.Builder
    lastPos := 0
    for _, vb := range validBlocks {
      if vb.startIdx > lastPos {
        remainingBuilder.WriteString(content[lastPos:vb.startIdx])
      }
      allCalls = append(allCalls, vb.calls...)
      lastPos = vb.endIdx
    }
    if lastPos < len(content) {
      remainingBuilder.WriteString(content[lastPos:])
    }

    for i := range allCalls {
      idx := i
      allCalls[i].Index = &idx
    }

    cleanRemaining := strings.TrimSpace(remainingBuilder.String())
    return allCalls, cleanRemaining, true
  }

  // Direct JSON fallback (only if every parsed call is present in allowedTools)
  trimmed := strings.TrimSpace(content)
  if (strings.HasPrefix(trimmed, "{") && strings.HasSuffix(trimmed, "}")) ||
    (strings.HasPrefix(trimmed, "[") && strings.HasSuffix(trimmed, "]")) ||
    (strings.HasPrefix(trimmed, "```json") && strings.HasSuffix(trimmed, "```")) {
    if directCalls, ok := parseToolCallsFromText(trimmed); ok {
      var validCalls []ToolCall
      for _, c := range directCalls {
        if allowedTools[c.Function.Name] {
          validCalls = append(validCalls, c)
        }
      }
      if len(validCalls) > 0 && len(validCalls) == len(directCalls) {
        for i := range validCalls {
          idx := i
          validCalls[i].Index = &idx
        }
        return validCalls, "", true
      }
    }
  }

  return nil, content, false
}

func ExtractThinking(content string) (string, string) {
  if strings.Contains(content, "<think>") && strings.Contains(content, "</think>") {
    start := strings.Index(content, "<think>")
    end := strings.Index(content, "</think>")
    if start < end {
      reasoning := content[start+len("<think>") : end]
      rem := content[:start] + content[end+len("</think>"):]
      rem = strings.TrimPrefix(rem, "\n\n")
      rem = strings.TrimPrefix(rem, "\n")
      return strings.TrimSpace(rem), strings.TrimSpace(reasoning)
    }
  }
  return content, ""
}

// ---------------------------------------------------------------------------
// Streaming Filters & Framing
// ---------------------------------------------------------------------------

type StreamThinkingFilter struct {
  inThinking bool
  buf        string
}

func NewStreamThinkingFilter() *StreamThinkingFilter {
  return &StreamThinkingFilter{}
}

func hasPrefixOf(target string, prefixes []string) int {
  for _, p := range prefixes {
    if strings.HasSuffix(target, p) {
      return len(p)
    }
  }
  return 0
}

func (f *StreamThinkingFilter) Feed(chunk string, onContent func(string), onReasoning func(string)) {
  f.buf += chunk
  thinkStartTag := "<think>"
  thinkEndTag := "</think>"

  thinkStartPrefixes := []string{"<", "<t", "<th", "<thi", "<thin", "<think"}
  thinkEndPrefixes := []string{"<", "</", "</t", "</th", "</thi", "</thin", "</think"}

  for len(f.buf) > 0 {
    if !f.inThinking {
      if idx := strings.Index(f.buf, thinkStartTag); idx != -1 {
        before := f.buf[:idx]
        if before != "" {
          onContent(before)
        }
        f.inThinking = true
        f.buf = f.buf[idx+len(thinkStartTag):]
      } else if matchLen := hasPrefixOf(f.buf, thinkStartPrefixes); matchLen > 0 {
        safe := f.buf[:len(f.buf)-matchLen]
        if safe != "" {
          onContent(safe)
        }
        f.buf = f.buf[len(f.buf)-matchLen:]
        break
      } else {
        onContent(f.buf)
        f.buf = ""
        break
      }
    } else {
      if idx := strings.Index(f.buf, thinkEndTag); idx != -1 {
        before := f.buf[:idx]
        if before != "" {
          onReasoning(before)
        }
        f.inThinking = false
        f.buf = f.buf[idx+len(thinkEndTag):]
        f.buf = strings.TrimPrefix(f.buf, "\n\n")
        f.buf = strings.TrimPrefix(f.buf, "\n")
      } else if matchLen := hasPrefixOf(f.buf, thinkEndPrefixes); matchLen > 0 {
        safe := f.buf[:len(f.buf)-matchLen]
        if safe != "" {
          onReasoning(safe)
        }
        f.buf = f.buf[len(f.buf)-matchLen:]
        break
      } else {
        onReasoning(f.buf)
        f.buf = ""
        break
      }
    }
  }
}

func (f *StreamThinkingFilter) Flush(onContent func(string), onReasoning func(string)) {
  if len(f.buf) > 0 {
    if f.inThinking {
      onReasoning(f.buf)
    } else {
      onContent(f.buf)
    }
    f.buf = ""
  }
}

type StreamToolCallFilter struct {
  allowedTools map[string]bool
  inToolCall   bool
  buf          string
  toolCallBuf  string
  toolIndex    int
  emittedCall  bool
  preambleBuf  string
}

func NewStreamToolCallFilter(allowedTools map[string]bool) *StreamToolCallFilter {
  return &StreamToolCallFilter{
    allowedTools: allowedTools,
  }
}

func (f *StreamToolCallFilter) appendPreambleOrContent(text string, onContent func(string)) {
  if text == "" {
    return
  }
  if f.emittedCall {
    onContent(text)
    return
  }
  if len(f.preambleBuf)+len(text) < 128 && !strings.Contains(f.preambleBuf+text, "\n\n") {
    f.preambleBuf += text
  } else {
    if f.preambleBuf != "" {
      onContent(f.preambleBuf)
      f.preambleBuf = ""
    }
    onContent(text)
  }
}

func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onReasoning func(string), onToolCall func(ToolCall)) {
  if len(f.allowedTools) == 0 {
    if chunk != "" {
      onContent(chunk)
    }
    return
  }

  f.buf += chunk
  toolStartTag := "<tool_call>"
  toolEndTag := "</tool_call>"

  startPrefixes := []string{"<", "<t", "<to", "<too", "<tool", "<tool_", "<tool_c", "<tool_ca", "<tool_cal", "<tool_call"}
  endPrefixes := []string{"<", "</", "</t", "</to", "</too", "</tool", "</tool_", "</tool_c", "</tool_ca", "</tool_cal", "</tool_call"}

  for len(f.buf) > 0 {
    if !f.inToolCall {
      if idx := strings.Index(f.buf, toolStartTag); idx != -1 {
        before := f.buf[:idx]
        f.inToolCall = true
        f.buf = f.buf[idx+len(toolStartTag):]
        f.preambleBuf += before
      } else if matchLen := hasPrefixOf(f.buf, startPrefixes); matchLen > 0 {
        safe := f.buf[:len(f.buf)-matchLen]
        f.buf = f.buf[len(f.buf)-matchLen:]
        f.appendPreambleOrContent(safe, onContent)
        break
      } else {
        safe := f.buf
        f.buf = ""
        f.appendPreambleOrContent(safe, onContent)
        break
      }
    } else {
      if idx := strings.Index(f.buf, toolEndTag); idx != -1 {
        f.toolCallBuf += f.buf[:idx]
        f.buf = f.buf[idx+len(toolEndTag):]
        f.inToolCall = false

        rawBlock := "<tool_call>" + f.toolCallBuf + "</tool_call>"
        tcs, ok := parseXMLToolCall(rawBlock)
        var validCalls []ToolCall
        if ok {
          for _, tc := range tcs {
            if f.allowedTools[tc.Function.Name] {
              validCalls = append(validCalls, tc)
            }
          }
        }

        if len(validCalls) > 0 {
          if strings.TrimSpace(f.preambleBuf) != "" {
            onReasoning(f.preambleBuf)
          }
          f.preambleBuf = ""
          for _, tc := range validCalls {
            idxCopy := f.toolIndex
            tc.Index = &idxCopy
            f.toolIndex++
            f.emittedCall = true
            onToolCall(tc)
          }
        } else {
          // Tool call invalid or function name not allowed - flush as content
          if f.preambleBuf != "" {
            onContent(f.preambleBuf)
            f.preambleBuf = ""
          }
          onContent(rawBlock)
        }
        f.toolCallBuf = ""
      } else if matchLen := hasPrefixOf(f.buf, endPrefixes); matchLen > 0 {
        safe := f.buf[:len(f.buf)-matchLen]
        f.toolCallBuf += safe
        f.buf = f.buf[len(f.buf)-matchLen:]
        break
      } else {
        f.toolCallBuf += f.buf
        f.buf = ""
        break
      }
    }
  }
}

func (f *StreamToolCallFilter) Flush(onContent func(string), onReasoning func(string), onToolCall func(ToolCall)) {
  if len(f.allowedTools) == 0 {
    if f.buf != "" {
      onContent(f.buf)
      f.buf = ""
    }
    return
  }

  if f.inToolCall && len(f.toolCallBuf) > 0 {
    rawBlock := "<tool_call>" + f.toolCallBuf + "</tool_call>"
    tcs, ok := parseXMLToolCall(rawBlock)
    var validCalls []ToolCall
    if ok {
      for _, tc := range tcs {
        if f.allowedTools[tc.Function.Name] {
          validCalls = append(validCalls, tc)
        }
      }
    }

    if len(validCalls) > 0 {
      if strings.TrimSpace(f.preambleBuf) != "" {
        onReasoning(f.preambleBuf)
      }
      f.preambleBuf = ""
      for _, tc := range validCalls {
        idxCopy := f.toolIndex
        tc.Index = &idxCopy
        f.toolIndex++
        f.emittedCall = true
        onToolCall(tc)
      }
    } else {
      if f.preambleBuf != "" {
        onContent(f.preambleBuf)
        f.preambleBuf = ""
      }
      onContent("<tool_call>" + f.toolCallBuf)
    }
    f.toolCallBuf = ""
    f.inToolCall = false
  }

  if f.emittedCall {
    if strings.TrimSpace(f.preambleBuf) != "" {
      onReasoning(f.preambleBuf)
    }
    f.preambleBuf = ""
  } else {
    if f.preambleBuf != "" {
      onContent(f.preambleBuf)
      f.preambleBuf = ""
    }
  }
  if len(f.buf) > 0 {
    onContent(f.buf)
    f.buf = ""
  }
}

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) {
  if text != "" {
    sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{ReasoningContent: text}, "")
  }
}

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

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

func (s *Streamer) Finish(reason string) {
  sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{}, reason)
  fmt.Fprintf(s.w, "data: [DONE]\n\n")
  if s.flusher != nil {
    s.flusher.Flush()
  }
}

func sendStreamDelta(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, delta StreamDelta, finishReason string) {
  chunk := StreamResponse{
    ID:      id,
    Object:  "chat.completion.chunk",
    Created: created,
    Model:   model,
    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()
  }
}

func WriteCompletionResponse(w http.ResponseWriter, completionID string, created int64, model string, content string, reasoning string, toolCalls []ToolCall, finishReason string) {
  if finishReason == "" {
    finishReason = "stop"
  }
  var contentVal interface{} = content
  if len(toolCalls) > 0 {
    finishReason = "tool_calls"
    if reasoning == "" && strings.TrimSpace(content) != "" {
      reasoning = strings.TrimSpace(content)
    }
    contentVal = nil
  }

  resp := ChatCompletionResponse{
    ID:      completionID,
    Object:  "chat.completion",
    Created: created,
    Model:   model,
    Choices: []ChatCompletionResponseChoice{
      {
        Index: 0,
        Message: ChatMessage{
          Role:             "assistant",
          Content:          contentVal,
          ReasoningContent: reasoning,
          ToolCalls:        toolCalls,
        },
        FinishReason: finishReason,
      },
    },
    Usage: Usage{
      PromptTokens:     0,
      CompletionTokens: 0,
      TotalTokens:      0,
    },
  }
  w.Header().Set("Content-Type", "application/json")
  json.NewEncoder(w).Encode(resp)
}

type APIErrorDetail struct {
  Message string  `json:"message"`
  Type    string  `json:"type"`
  Param   *string `json:"param"`
  Code    int     `json:"code"`
}

type APIErrorResponse struct {
  Error APIErrorDetail `json:"error"`
}

func WriteAPIError(w http.ResponseWriter, status int, message string, errType string, code int) {
  w.Header().Set("Content-Type", "application/json")
  w.WriteHeader(status)
  resp := APIErrorResponse{
    Error: APIErrorDetail{
      Message: message,
      Type:    errType,
      Code:    code,
    },
  }
  _ = json.NewEncoder(w).Encode(resp)
}

// ---------------------------------------------------------------------------
// Gateway Controller & Request Orchestration
// ---------------------------------------------------------------------------

type GroqqerGateway struct {
  targetURL    string
  proxyURL     string
  userAgent    string
  defaultModel string
  port         int
  models       []ModelItem
  modelsMu     sync.RWMutex
}

func DefaultFallbackModels() []string {
  return []string{
    "qwen/qwen3.6-27b",
    "qwen/qwen3.8-27b",
    "openai/gpt-oss-120b",
    "openai/gpt-oss-20b",
    "groq/compound",
    "groq/compound-mini",
    "allam-2-7b",
    "llama-3.3-70b-versatile",
    "llama-3.1-8b-instant",
    "meta-llama/llama-4-scout-17b-16e-instruct",
  }
}

func (g *GroqqerGateway) RefreshModels() {
  ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
  defer cancel()

  sess, err := OpenStreamlitSession(ctx, g.targetURL, g.proxyURL, g.userAgent)
  if err != nil {
    log.Printf("warning: model discovery session failed: %v", err)
    if len(g.GetModels()) == 0 {
      g.setFallbackModels()
    }
    return
  }
  defer sess.Close()

  now := time.Now().Unix()
  var items []ModelItem
  for _, opt := range sess.options {
    rawOpt := opt
    if idx := strings.Index(rawOpt, " ("); idx != -1 {
      rawOpt = strings.TrimSpace(rawOpt[:idx])
    }
    if rawOpt != "" {
      items = append(items, ModelItem{
        ID:      rawOpt,
        Object:  "model",
        Created: now,
        OwnedBy: "groq",
      })
    }
  }

  if len(items) == 0 {
    g.setFallbackModels()
    return
  }

  g.modelsMu.Lock()
  g.models = items
  g.modelsMu.Unlock()
}

func (g *GroqqerGateway) setFallbackModels() {
  g.modelsMu.Lock()
  defer g.modelsMu.Unlock()
  now := time.Now().Unix()
  var items []ModelItem
  for _, id := range DefaultFallbackModels() {
    items = append(items, ModelItem{
      ID:      id,
      Object:  "model",
      Created: now,
      OwnedBy: "groq",
    })
  }
  g.models = items
}

func (g *GroqqerGateway) GetModels() []ModelItem {
  g.modelsMu.RLock()
  defer g.modelsMu.RUnlock()
  if len(g.models) == 0 {
    now := time.Now().Unix()
    var items []ModelItem
    for _, id := range DefaultFallbackModels() {
      items = append(items, ModelItem{
        ID:      id,
        Object:  "model",
        Created: now,
        OwnedBy: "groq",
      })
    }
    return items
  }
  out := make([]ModelItem, len(g.models))
  copy(out, g.models)
  return out
}

func (g *GroqqerGateway) MatchModel(requestedModel string) string {
  if requestedModel == "" {
    return g.defaultModel
  }
  models := g.GetModels()
  for _, m := range models {
    if strings.EqualFold(m.ID, requestedModel) {
      return m.ID
    }
  }
  // Try suffix match e.g. "qwen3.6-27b" for "qwen/qwen3.6-27b"
  cleanReq := strings.ToLower(strings.TrimSpace(requestedModel))
  if strings.Contains(cleanReq, "/") {
    parts := strings.Split(cleanReq, "/")
    cleanReq = parts[len(parts)-1]
  }
  for _, m := range models {
    if strings.Contains(strings.ToLower(m.ID), cleanReq) {
      return m.ID
    }
  }
  return requestedModel
}

func (g *GroqqerGateway) HandleModels(w http.ResponseWriter, r *http.Request) {
  w.Header().Set("Access-Control-Allow-Origin", "*")
  w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
  w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
  if r.Method == "OPTIONS" {
    w.WriteHeader(http.StatusOK)
    return
  }

  models := g.GetModels()
  resp := ModelsResponse{
    Object: "list",
    Data:   models,
  }
  w.Header().Set("Content-Type", "application/json")
  json.NewEncoder(w).Encode(resp)
}

func (g *GroqqerGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Request) {
  w.Header().Set("Access-Control-Allow-Origin", "*")
  w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
  w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
  if r.Method == "OPTIONS" {
    w.WriteHeader(http.StatusOK)
    return
  }

  if r.Method != "POST" {
    WriteAPIError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", 405)
    return
  }

  var req ChatCompletionRequest
  if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
    WriteAPIError(w, http.StatusBadRequest, "Invalid JSON payload: "+err.Error(), "invalid_request_error", 400)
    return
  }

  if len(req.Messages) == 0 {
    WriteAPIError(w, http.StatusBadRequest, "messages array must not be empty", "invalid_request_error", 400)
    return
  }

  requestedModel := g.MatchModel(req.Model)
  promptText := FormatPrompt(req)
  completionID := "chatcmpl-" + GenerateUUID()
  created := time.Now().Unix()
  allowedTools := GetAllowedToolNames(req.Tools, req.ToolChoice)

  log.Printf("Handling completion: model=%s, stream=%v, messages=%d", requestedModel, req.Stream, len(req.Messages))

  ctx, cancel := context.WithTimeout(r.Context(), 180*time.Second)
  defer cancel()

  start := time.Now()
  maxAttempts := 3

  for attempt := 0; attempt < maxAttempts; attempt++ {
    sess, err := OpenStreamlitSession(ctx, g.targetURL, g.proxyURL, g.userAgent)
    if err != nil {
      log.Printf("session connection error (attempt %d): %v", attempt+1, err)
      if attempt < maxAttempts-1 {
        time.Sleep(1 * time.Second)
        continue
      }
      WriteAPIError(w, http.StatusBadGateway, "Failed to connect to target space: "+err.Error(), "api_error", 502)
      return
    }

    if err := sess.SwitchModel(requestedModel); err != nil {
      log.Printf("model switch warning: %v", err)
    }

    effectiveModel := sess.activeModel
    if effectiveModel == "" {
      effectiveModel = requestedModel
    }

    if req.Stream {
      flusher, ok := w.(http.Flusher)
      if !ok {
        sess.Close()
        WriteAPIError(w, http.StatusInternalServerError, "Streaming unsupported", "api_error", 500)
        return
      }

      var streamer *Streamer
      var thinkingFilter *StreamThinkingFilter
      var toolFilter *StreamToolCallFilter
      var streamStarted bool
      var alertMsg string

      initStreamer := func() {
        if !streamStarted {
          streamer = NewStreamer(w, flusher, completionID, created, effectiveModel)
          streamer.Role()
          thinkingFilter = NewStreamThinkingFilter()
          toolFilter = NewStreamToolCallFilter(allowedTools)
          streamStarted = true
        }
      }

      err = sess.SubmitPrompt(promptText,
        func(delta string) {
          initStreamer()
          thinkingFilter.Feed(delta,
            func(cleanContent string) {
              toolFilter.Feed(cleanContent,
                func(userText string) {
                  streamer.Content(userText)
                },
                func(reasoning string) {
                  streamer.Reasoning(reasoning)
                },
                func(tc ToolCall) {
                  streamer.ToolCall(tc)
                },
              )
            },
            func(reasoning string) {
              streamer.Reasoning(reasoning)
            },
          )
        },
        func(alert string) {
          alertMsg = alert
        },
      )
      sess.Close()

      if err != nil {
        var alertErr *GatewayAlertError
        if errors.As(err, &alertErr) && !streamStarted && alertErr.IsRateLimit {
          if attempt < maxAttempts-1 {
            waitSec := alertErr.RetryAfter
            if waitSec <= 0 {
              waitSec = 6
            }
            if waitSec > 20 {
              waitSec = 20
            }
            log.Printf("Rate limit hit on streaming %s. Waiting %d seconds before retry (attempt %d)...", effectiveModel, waitSec, attempt+1)
            select {
            case <-time.After(time.Duration(waitSec)*time.Second + 500*time.Millisecond):
              continue
            case <-ctx.Done():
              WriteAPIError(w, http.StatusTooManyRequests, "Rate limit retry cancelled: context deadline exceeded", "rate_limit_error", 429)
              return
            }
          }
          w.Header().Set("Retry-After", strconv.Itoa(alertErr.RetryAfter))
          WriteAPIError(w, http.StatusTooManyRequests, alertErr.Alert, "rate_limit_error", 429)
          return
        }

        if !streamStarted {
          if errors.As(err, &alertErr) {
            lower := strings.ToLower(alertErr.Alert)
            if strings.Contains(lower, "too large") || strings.Contains(lower, "cannot fit") {
              WriteAPIError(w, http.StatusBadRequest, alertErr.Alert, "invalid_request_error", 400)
              return
            }
            WriteAPIError(w, http.StatusBadGateway, alertErr.Alert, "api_error", 502)
            return
          }
          WriteAPIError(w, http.StatusInternalServerError, "Streaming error: "+err.Error(), "api_error", 500)
          return
        }

        if alertMsg != "" {
          streamer.Content("\n\n" + alertMsg)
        }
      }

      if streamStarted {
        finishReason := "stop"
        thinkingFilter.Flush(
          func(cleanContent string) {
            toolFilter.Feed(cleanContent,
              func(userText string) {
                streamer.Content(userText)
              },
              func(reasoning string) {
                streamer.Reasoning(reasoning)
              },
              func(tc ToolCall) {
                streamer.ToolCall(tc)
              },
            )
          },
          func(reasoning string) {
            streamer.Reasoning(reasoning)
          },
        )

        toolFilter.Flush(
          func(userText string) {
            streamer.Content(userText)
          },
          func(reasoning string) {
            streamer.Reasoning(reasoning)
          },
          func(tc ToolCall) {
            streamer.ToolCall(tc)
          },
        )

        if toolFilter.emittedCall {
          finishReason = "tool_calls"
        }

        streamer.Finish(finishReason)
        log.Printf("Streaming completion finished for %s in %v", effectiveModel, time.Since(start))
        return
      }

      WriteAPIError(w, http.StatusInternalServerError, "Empty response from space", "api_error", 500)
      return
    }

    // Non-streaming completion
    var fullBuilder strings.Builder
    var alertMsg string

    err = sess.SubmitPrompt(promptText,
      func(delta string) {
        fullBuilder.WriteString(delta)
      },
      func(alert string) {
        alertMsg = alert
      },
    )
    sess.Close()
    _ = alertMsg

    if err != nil {
      var alertErr *GatewayAlertError
      if errors.As(err, &alertErr) && alertErr.IsRateLimit {
        if attempt < maxAttempts-1 {
          waitSec := alertErr.RetryAfter
          if waitSec <= 0 {
            waitSec = 6
          }
          if waitSec > 20 {
            waitSec = 20
          }
          log.Printf("Rate limit hit on %s. Waiting %d seconds before retry (attempt %d)...", effectiveModel, waitSec, attempt+1)
          select {
          case <-time.After(time.Duration(waitSec)*time.Second + 500*time.Millisecond):
            continue
          case <-ctx.Done():
            WriteAPIError(w, http.StatusTooManyRequests, "Rate limit retry cancelled: context deadline exceeded", "rate_limit_error", 429)
            return
          }
        }
        w.Header().Set("Retry-After", strconv.Itoa(alertErr.RetryAfter))
        WriteAPIError(w, http.StatusTooManyRequests, alertErr.Alert, "rate_limit_error", 429)
        return
      }

      if errors.As(err, &alertErr) {
        lower := strings.ToLower(alertErr.Alert)
        if strings.Contains(lower, "too large") || strings.Contains(lower, "cannot fit") {
          WriteAPIError(w, http.StatusBadRequest, alertErr.Alert, "invalid_request_error", 400)
          return
        }
        WriteAPIError(w, http.StatusBadGateway, alertErr.Alert, "api_error", 502)
        return
      }

      WriteAPIError(w, http.StatusInternalServerError, "Prompt execution error: "+err.Error(), "api_error", 500)
      return
    }

    fullText := fullBuilder.String()
    if strings.TrimSpace(fullText) == "" {
      WriteAPIError(w, http.StatusInternalServerError, "The model returned an empty response", "api_error", 500)
      return
    }

    cleanContent, reasoningContent := ExtractThinking(fullText)
    toolCalls, remainingContent, hasTools := DetectToolCalls(cleanContent, allowedTools)

    finishReason := "stop"
    if hasTools {
      finishReason = "tool_calls"
      cleanContent = remainingContent
    }

    WriteCompletionResponse(w, completionID, created, effectiveModel, cleanContent, reasoningContent, toolCalls, finishReason)
    log.Printf("Completion finished for %s in %v", effectiveModel, time.Since(start))
    return
  }
}

// ---------------------------------------------------------------------------
// Main Entrypoint & Signal Handling
// ---------------------------------------------------------------------------

func main() {
  port := flag.Int("port", 8080, "HTTP server listening port")
  targetURL := flag.String("target", DefaultTargetURL, "Groq Streamlit space URL")
  defaultModel := flag.String("default-model", DefaultModel, "Default model ID if unspecified")
  userAgent := flag.String("user-agent", DefaultUserAgent, "User-Agent string")
  flag.StringVar(userAgent, "ua", DefaultUserAgent, "User-Agent string (alias)")

  socksProxy := flag.String("socks", "", "SOCKS5 proxy (e.g. socks5://127.0.0.1:1080)")
  flag.StringVar(socksProxy, "proxy", "", "Proxy address alias")
  flag.StringVar(socksProxy, "socks5", "", "SOCKS5 proxy alias")

  // Retain flags for backward compatibility (browser/Xvfb are no longer required)
  _ = flag.String("browser", "", "Ignored: browser is no longer required (headless WebSocket mode)")
  _ = flag.Bool("xvfb", false, "Ignored: Xvfb is no longer required")
  _ = flag.Bool("no-xvfb", false, "Ignored: Xvfb is no longer required")
  _ = flag.Bool("headless", true, "Ignored: always runs headless via direct WebSocket")

  flag.Parse()

  effectiveProxy := *socksProxy
  if effectiveProxy == "" {
    for _, envVar := range []string{"ALL_PROXY", "all_proxy", "SOCKS5_PROXY", "socks5_proxy", "SOCKS_PROXY", "socks_proxy"} {
      if val := os.Getenv(envVar); val != "" {
        effectiveProxy = val
        break
      }
    }
  }

  log.Printf("Initializing groqqer gateway (pure Go WebSocket client, zero browser dependency)...")

  gateway := &GroqqerGateway{
    targetURL:    *targetURL,
    proxyURL:     effectiveProxy,
    userAgent:    *userAgent,
    defaultModel: *defaultModel,
    port:         *port,
  }

  // Initial discovery of live models in background
  go func() {
    log.Println("Discovering available models from Streamlit space...")
    gateway.RefreshModels()
    models := gateway.GetModels()
    log.Printf("Ready with %d active models. Default: %s", len(models), gateway.defaultModel)
  }()

  mux := http.NewServeMux()
  mux.HandleFunc("/v1/models", gateway.HandleModels)
  mux.HandleFunc("/models", gateway.HandleModels)
  mux.HandleFunc("/v1/chat/completions", gateway.HandleChatCompletions)
  mux.HandleFunc("/chat/completions", gateway.HandleChatCompletions)

  // Root status handler
  mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/" {
      http.NotFound(w, r)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]interface{}{
      "service":   "groqqer",
      "mode":      "headless-websocket-protobuf",
      "status":    "running",
      "target":    *targetURL,
      "default":   *defaultModel,
      "models":    "/v1/models",
      "endpoints": []string{"/v1/chat/completions", "/v1/models"},
    })
  })

  server := &http.Server{
    Addr:    fmt.Sprintf(":%d", *port),
    Handler: mux,
  }

  // Graceful shutdown handling
  sigChan := make(chan os.Signal, 1)
  signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

  go func() {
    <-sigChan
    log.Println("Shutting down groqqer...")
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    _ = server.Shutdown(ctx)
    os.Exit(0)
  }()

  log.Printf("groqqer gateway listening on http://127.0.0.1:%d", *port)
  if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
    log.Fatalf("server error: %v", err)
  }
}