"""Reliable ZeroGPU backend for the local OpenAI-compatible proxy."""

from __future__ import annotations

import json
import os
import time
import uuid
import traceback
import threading
from typing import Any

os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")

import gradio as gr
from gradio.context import LocalContext
import spaces
import torch
from fastapi import HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field, ValidationError
from starlette.concurrency import run_in_threadpool
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    StoppingCriteria,
    StoppingCriteriaList,
)
from generation import (
    gpu_duration_seconds,
    head_tail_token_counts,
    merge_eos_token_ids,
)
from openai_compat import (
    analyze_tool_flow,
    indexed_tool_calls,
    normalize_tools,
    resolve_tool_choice,
    select_tools,
    tool_choice_instruction,
    tool_protocol_instruction,
    tool_names,
)
from openclaude_compat import (
    TOOL_PROTOCOL_MARKER,
    add_system_instruction,
    has_tool_protocol,
    normalize_openclaude_messages,
)
from tool_calls import (
    extract_tool_calls,
    has_complete_tool_call,
    recover_forced_tool_call,
)
from web_search import SearchUnavailable, search_web


# O Titã: Qwen2.5-Coder-32B nativamente quantizado em 4-bits (AWQ)
MODEL = os.getenv(
    "MODEL",
    os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"),
)

NATIVE_CONTEXT_TOKENS = 32768
MAX_SUPPORTED_CONTEXT_TOKENS = 131072
MAX_CONTEXT_TOKENS = int(os.getenv("MAX_CONTEXT_TOKENS", str(MAX_SUPPORTED_CONTEXT_TOKENS)))
if not 1 <= MAX_CONTEXT_TOKENS <= MAX_SUPPORTED_CONTEXT_TOKENS:
    raise RuntimeError(
        f"MAX_CONTEXT_TOKENS must be between 1 and {MAX_SUPPORTED_CONTEXT_TOKENS}; "
        f"got {MAX_CONTEXT_TOKENS}"
    )
YARN_ENABLED = MAX_CONTEXT_TOKENS > NATIVE_CONTEXT_TOKENS
YARN_FACTOR = MAX_CONTEXT_TOKENS / NATIVE_CONTEXT_TOKENS
ZERO_GPU_SIZE = "xlarge" if YARN_ENABLED else "large"
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "2048"))
MAX_TOOL_CALL_TOKENS = int(os.getenv("MAX_TOOL_CALL_TOKENS", "2048"))
DEFAULT_TEMPERATURE = float(os.getenv("DEFAULT_TEMPERATURE", "0.0"))
MAX_TEMPERATURE = float(os.getenv("MAX_TEMPERATURE", "0.2"))
TOOL_TEMPERATURE = 0.0
PRESERVED_PREFIX_TOKENS = int(os.getenv("PRESERVED_PREFIX_TOKENS", "4096"))
MODEL_ALIASES = tuple(dict.fromkeys((MODEL, "qwen2.5-coder-32b", "qwen-coder")))

tokenizer = AutoTokenizer.from_pretrained(MODEL)

# O ZeroGPU só anexa uma GPU real dentro de funções decoradas com
# @spaces.GPU; no escopo do módulo (startup) não existe CUDA de verdade,
# apenas uma emulação que aceita `.to("cuda")`/`device_map="auto"` como
# simples posicionamento de tensores. O carregamento deste modelo AWQ,
# porém, dispara o kernel Marlin (`awq_marlin_repack`) de forma síncrona
# dentro do próprio from_pretrained — isso é execução real de kernel CUDA,
# não posicionamento, e por isso não existe backend CPU para ele (era
# exatamente esse o erro do seu log). Por isso o carregamento precisa ser
# adiado para dentro de `gerar`, a única função com GPU real anexada.
model: AutoModelForCausalLM | None = None
_MODEL_LOAD_LOCK = threading.Lock()
# Custom HTTP routes bypass Gradio's event queue. Serialize calls *before*
# entering @spaces.GPU so one global AWQ model is never generated from by two
# request threads at the same time. This also prevents concurrent cold loads.
_GENERATION_LOCK = threading.Lock()


def _ensure_model_loaded() -> None:
    """Load AWQ once per worker and never expose a half-initialized model."""
    global model
    if model is not None:
        return

    with _MODEL_LOAD_LOCK:
        if model is not None:
            return
        print(f"Loading {MODEL} on ZeroGPU (NATIVE AWQ)...", flush=True)
        model_kwargs: dict[str, Any] = {
            "dtype": "auto",
            "device_map": "auto",
            "low_cpu_mem_usage": True,
        }
        if YARN_ENABLED:
            # Qwen2.5 is natively configured for 32K. The official model card
            # recommends YaRN for longer contexts, up to 131,072 tokens.
            # transformers==5.14.1 uses Qwen2Config.rope_parameters (not the
            # pre-5.x rope_scaling name). YaRN validation requires the original
            # pretrained window; rope_theta mirrors this checkpoint's config.
            model_kwargs["rope_parameters"] = {
                "rope_type": "yarn",
                "factor": float(YARN_FACTOR),
                "original_max_position_embeddings": NATIVE_CONTEXT_TOKENS,
                "rope_theta": 1_000_000.0,
            }
            model_kwargs["max_position_embeddings"] = MAX_CONTEXT_TOKENS
        candidate = AutoModelForCausalLM.from_pretrained(
            MODEL,
            **model_kwargs,
        )
        candidate.eval()
        model = candidate
        print(f"Model ready on {next(model.parameters()).device}", flush=True)


def _bounded_output_tokens(value: float) -> int:
    try:
        requested = int(value)
    except (TypeError, ValueError):
        requested = MAX_NEW_TOKENS
    return max(1, min(requested, MAX_NEW_TOKENS))


# Buffer para cobrir a compilação JIT do kernel Marlin + carregamento dos
# pesos quando `gerar` cai num worker "frio" (sem o modelo em memória).
# É uma estimativa (baseada nos ~99s de compilação que aparecem no seu log);
# meça o cold start real do seu Space e ajuste. Confira também o teto de
# duração por chamada da sua tier em
# https://huggingface.co/docs/hub/spaces-zerogpu antes de subir esse valor —
# se o teto for menor que isso, a chamada falha com "illegal duration".
COLD_START_BUFFER_SECONDS = 180


def _gpu_duration(
    messages_json: str,
    __: float,
    max_new_tokens: float,
    *tool_arguments: object,
) -> int:
    output_tokens = _bounded_output_tokens(max_new_tokens)
    tool_characters = sum(
        len(value) for value in tool_arguments if isinstance(value, str)
    )
    duration = gpu_duration_seconds(
        len(messages_json) + tool_characters,
        output_tokens,
        MAX_CONTEXT_TOKENS,
    )
    cold_start_buffer = COLD_START_BUFFER_SECONDS if model is None else 0
    return duration + cold_start_buffer


def _tool_protocol_active(messages: list[object]) -> bool:
    return any(
        isinstance(message, dict)
        and isinstance(message.get("content"), str)
        and TOOL_PROTOCOL_MARKER in message["content"]
        for message in messages
    )


def _native_tools(raw_tools: object) -> list[dict[str, Any]]:
    return normalize_tools(raw_tools)


def _render_prompt(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
) -> str:
    """Render the exact Qwen prompt and never silently discard requested tools."""
    template_kwargs: dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": True,
    }
    if tools:
        template_kwargs["tools"] = tools

    try:
        return tokenizer.apply_chat_template(messages, **template_kwargs)
    except Exception as template_error:
        if tools:
            raise RuntimeError(
                "Qwen chat template failed while tools were enabled; refusing "
                "to continue with a tool-less prompt"
            ) from template_error
        raise


def _prompt_input_ids(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
) -> list[int]:
    prompt = _render_prompt(messages, tools)
    return tokenizer(
        prompt,
        add_special_tokens=False,
        truncation=False,
    )["input_ids"]


def _trim_oldest_turn(messages: list[dict[str, Any]]) -> list[dict[str, Any]] | None:
    """Drop one old conversation turn while retaining every system instruction."""
    user_indexes = [
        index
        for index, message in enumerate(messages)
        if str(message.get("role", "")).casefold() == "user"
    ]
    if len(user_indexes) >= 2:
        cutoff = user_indexes[1]
        return [
            message
            for index, message in enumerate(messages)
            if index >= cutoff
            or str(message.get("role", "")).casefold() == "system"
        ]

    if user_indexes and user_indexes[0] > 0:
        cutoff = user_indexes[0]
        trimmed = [
            message
            for index, message in enumerate(messages)
            if index >= cutoff
            or str(message.get("role", "")).casefold() == "system"
        ]
        return trimmed if trimmed != messages else None
    return None


CONTEXT_TRUNCATION_MARKER = "\n...[older/oversized content truncated to fit context]...\n"


def _truncate_text_to_tokens(text: str, target_tokens: int) -> str:
    """Shrink message *content* while preserving its surrounding chat syntax."""
    ids = tokenizer(text, add_special_tokens=False, truncation=False)["input_ids"]
    target = max(1, int(target_tokens))
    if len(ids) <= target:
        return text

    marker_ids = tokenizer(
        CONTEXT_TRUNCATION_MARKER,
        add_special_tokens=False,
        truncation=False,
    )["input_ids"]
    payload_budget = max(1, target - len(marker_ids))
    head = max(1, payload_budget // 2)
    tail = max(0, payload_budget - head)
    head_text = tokenizer.decode(ids[:head], skip_special_tokens=False)
    tail_text = (
        tokenizer.decode(ids[-tail:], skip_special_tokens=False)
        if tail
        else ""
    )
    return head_text + CONTEXT_TRUNCATION_MARKER + tail_text


def _fit_messages_to_context(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    output_tokens: int,
) -> list[dict[str, Any]]:
    """Fit context without ever slicing Qwen's rendered tool catalog.

    Old complete turns are removed first. If a tool-enabled request is still too
    large, textual message contents are reduced in-place at message boundaries,
    preserving the Qwen `<tools>` schema and all role/tool-call wrappers.
    """
    input_budget = max(1, MAX_CONTEXT_TOKENS - output_tokens)
    fitted = [dict(message) for message in messages]

    while len(_prompt_input_ids(fitted, tools)) > input_budget:
        trimmed = _trim_oldest_turn(fitted)
        if trimmed is None or trimmed == fitted:
            break
        fitted = trimmed

    if not tools:
        return fitted

    latest_user_index = max(
        (
            index
            for index, message in enumerate(fitted)
            if str(message.get("role", "")).casefold() == "user"
        ),
        default=-1,
    )

    # Oversized OpenClaude system prompts and tool results are data inside chat
    # messages. Compact those before considering any raw token slicing. Keep the
    # backend tool protocol itself intact because it defines the wire contract.
    for _ in range(max(8, len(fitted) * 4)):
        current_length = len(_prompt_input_ids(fitted, tools))
        if current_length <= input_budget:
            return fitted

        excess = current_length - input_budget
        candidates: list[tuple[int, int, int]] = []
        for index, message in enumerate(fitted):
            content = message.get("content")
            if not isinstance(content, str) or not content:
                continue
            if TOOL_PROTOCOL_MARKER in content:
                continue
            role = str(message.get("role", "")).casefold()
            minimum = 768 if index == latest_user_index else (512 if role in {"system", "tool"} else 256)
            token_length = len(
                tokenizer(content, add_special_tokens=False, truncation=False)["input_ids"]
            )
            if token_length > minimum:
                candidates.append((token_length, index, minimum))

        if not candidates:
            break

        token_length, index, minimum = max(candidates)
        target = max(minimum, token_length - excess - 64)
        if target >= token_length:
            target = max(minimum, token_length // 2)
        original = str(fitted[index]["content"])
        shortened = _truncate_text_to_tokens(original, target)
        if shortened == original:
            break
        fitted[index] = {**fitted[index], "content": shortened}

    final_length = len(_prompt_input_ids(fitted, tools))
    if final_length > input_budget:
        raise ValueError(
            "tool-enabled prompt exceeds the configured context window even "
            "after whole-turn and message-content compaction; refusing to slice "
            "the Qwen tool schema"
        )
    return fitted


def _prompt_token_count(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    output_tokens: int,
) -> int:
    """Count the prompt tokens that survive the same context bound as generation."""
    fitted = _fit_messages_to_context(messages, tools, output_tokens)
    encoded = _prompt_input_ids(fitted, tools)
    input_budget = max(1, MAX_CONTEXT_TOKENS - output_tokens)
    return min(len(encoded), input_budget)


def _completion_token_count(text: str) -> int:
    """Count visible generated tokens for OpenAI-compatible usage reporting."""
    return len(
        tokenizer(
            text,
            add_special_tokens=False,
            truncation=False,
        )["input_ids"]
    )


class StopAfterToolCall(StoppingCriteria):
    def __init__(self, prompt_length: int, allowed_names: set[str]) -> None:
        self.prompt_length = prompt_length
        self.allowed_names = set(allowed_names)

    def __call__(self, input_ids, scores, **_: object):
        completed = []
        for sequence in input_ids:
            generated = sequence[self.prompt_length :]
            text = tokenizer.decode(generated, skip_special_tokens=False)
            completed.append(has_complete_tool_call(text, self.allowed_names))
        return torch.tensor(completed, dtype=torch.bool, device=input_ids.device)


@spaces.GPU(duration=_gpu_duration, size=ZERO_GPU_SIZE)
def gerar(
    messages_json: str,
    temperature: float,
    max_new_tokens: float,
    tools_json: str = "[]",
    stop_after_first_tool: bool = True,
) -> str:
    _ensure_model_loaded()
    messages = json.loads(messages_json)
    if not isinstance(messages, list):
        raise ValueError("messages_json must contain a JSON list")
    try:
        tools = _native_tools(json.loads(tools_json))
    except (TypeError, ValueError, json.JSONDecodeError):
        tools = []
    if not isinstance(tools, list):
        tools = []

    output_tokens = _bounded_output_tokens(max_new_tokens)
    # A stopping criterion only makes sense when the current request actually
    # advertises functions. A stale protocol marker without a live tool catalog
    # must never make ordinary text stop on tool-like syntax.
    tool_mode = bool(tools)
    messages = _fit_messages_to_context(messages, tools, output_tokens)
    prompt = _render_prompt(messages, tools)

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
        truncation=False,
    )
    input_budget = max(1, MAX_CONTEXT_TOKENS - output_tokens)
    input_length = inputs["input_ids"].shape[1]
    
    if input_length > input_budget:
        if tools:
            raise ValueError(
                "tool-enabled prompt still exceeds context after safe compaction; "
                "refusing to slice the rendered <tools> catalog"
            )
        head_tokens, tail_tokens = head_tail_token_counts(
            input_length,
            input_budget,
            PRESERVED_PREFIX_TOKENS,
        )
        for key, value in inputs.items():
            if (
                isinstance(value, torch.Tensor)
                and value.ndim == 2
                and value.shape[1] == input_length
            ):
                parts = []
                if head_tokens:
                    parts.append(value[:, :head_tokens])
                if tail_tokens:
                    parts.append(value[:, -tail_tokens:])
                inputs[key] = torch.cat(parts, dim=1)
                
    inputs = inputs.to("cuda")

    print(
        f"Generation started: input_tokens={inputs['input_ids'].shape[1]} "
        f"max_new_tokens={output_tokens} tool_mode={tool_mode}",
        flush=True,
    )

    eos_token_ids = merge_eos_token_ids(
        model.generation_config.eos_token_id,
        tokenizer.eos_token_id,
    )
    
    generation_kwargs = {
        "max_new_tokens": output_tokens,
        "do_sample": float(temperature) > 0,
        "pad_token_id": (
            tokenizer.pad_token_id
            if tokenizer.pad_token_id is not None
            else tokenizer.eos_token_id
        ),
    }
    if eos_token_ids is not None:
        generation_kwargs["eos_token_id"] = eos_token_ids
        
    if generation_kwargs["do_sample"]:
        generation_kwargs["temperature"] = max(0.01, float(temperature))
        generation_kwargs["top_p"] = 0.8
        generation_kwargs["top_k"] = 20
        generation_kwargs["repetition_penalty"] = 1.05
        
    if tool_mode and stop_after_first_tool:
        generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
            [StopAfterToolCall(inputs["input_ids"].shape[1], tool_names(tools))]
        )

    with torch.inference_mode():
        output = model.generate(**inputs, **generation_kwargs)

    generated = output[0][inputs["input_ids"].shape[1] :]
    response = tokenizer.decode(generated, skip_special_tokens=True).strip()
    print(f"Generation completed: output_tokens={generated.shape[0]}", flush=True)
    return response


class ChatCompletionRequest(BaseModel):
    model: str = MODEL
    messages: list[dict[str, Any]] = Field(min_length=1)
    temperature: float = Field(default=DEFAULT_TEMPERATURE, ge=0.0)
    max_tokens: int | None = Field(default=None, ge=1)
    max_completion_tokens: int | None = Field(default=None, ge=1)
    stream: bool = False
    tools: list[dict[str, Any]] | None = None
    tool_choice: Any = None
    parallel_tool_calls: bool | None = None
    stream_options: dict[str, Any] | None = None


def _completion_payload(request: ChatCompletionRequest) -> dict[str, Any]:
    if request.model not in MODEL_ALIASES:
        raise HTTPException(status_code=404, detail=f"Model not available: {request.model}")

    already_adapted = has_tool_protocol(request.messages)
    flow_state = analyze_tool_flow(request.messages, request.tools or [])
    requested_mode = (
        request.tool_choice.casefold()
        if isinstance(request.tool_choice, str)
        else None
    )
    state_controls_choice = request.tool_choice is None or requested_mode in {"auto", "required"}
    effective_choice = resolve_tool_choice(request.tool_choice, flow_state)
    
    try:
        effective_tools, tool_mode = select_tools(
            request.tools or [], effective_choice
        )
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error

    instructions = [
        instruction
        for instruction in (
            (
                tool_protocol_instruction(
                    effective_tools,
                    parallel_tool_calls=request.parallel_tool_calls is True,
                )
                if effective_tools and not has_tool_protocol(request.messages)
                else None
            ),
            tool_choice_instruction(tool_mode, effective_tools),
            (
                flow_state.instruction
                if (
                    state_controls_choice
                    and not already_adapted
                    and not (
                        requested_mode == "required"
                        and flow_state.can_finalize
                        and not flow_state.requires_tool
                    )
                )
                else None
            ),
        )
        if instruction
    ]
    instruction = "\n\n".join(instructions) if instructions else None
    if request.max_completion_tokens is not None:
        max_tokens = request.max_completion_tokens
    elif request.max_tokens is not None:
        max_tokens = request.max_tokens
    else:
        max_tokens = MAX_NEW_TOKENS
    
    if effective_tools:
        max_tokens = min(max_tokens, MAX_TOOL_CALL_TOKENS)
    temperature = min(max(float(request.temperature), 0.0), MAX_TEMPERATURE)
    if tool_mode in {"required", "forced"}:
        # Tool JSON is a protocol surface, not creative prose. Greedy decoding
        # makes required/forced calls maximally reproducible and reduces malformed
        # argument objects on small/quantized models.
        temperature = TOOL_TEMPERATURE
    
    try:
        normalized_messages = (
            [dict(message) for message in request.messages]
            if already_adapted
            else normalize_openclaude_messages(request.messages)
        )
        prompt_messages = add_system_instruction(
            normalized_messages,
            instruction,
        )
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    
    bounded_max_tokens = _bounded_output_tokens(max_tokens)
    try:
        prompt_tokens = _prompt_token_count(
            prompt_messages,
            effective_tools,
            bounded_max_tokens,
        )
    except ValueError as error:
        raise HTTPException(status_code=413, detail=str(error)) from error
    with _GENERATION_LOCK:
        text = gerar(
            json.dumps(prompt_messages),
            temperature,
            bounded_max_tokens,
            json.dumps(effective_tools, ensure_ascii=False),
            request.parallel_tool_calls is not True,
        )
    completion_tokens = _completion_token_count(text)
    
    if effective_tools:
        tool_calls, content = extract_tool_calls(text, tool_names(effective_tools))
        if (
            not tool_calls
            and tool_mode in {"forced", "required"}
            and len(effective_tools) == 1
        ):
            recovered = recover_forced_tool_call(
                text,
                effective_tools[0]["function"]["name"],
            )
            if recovered is not None:
                tool_calls, content = [recovered], ""
        if request.parallel_tool_calls is False:
            tool_calls = tool_calls[:1]
    else:
        tool_calls, content = [], text
        
    # If Qwen emitted a complete tool-shaped payload but it did not validate
    # against the advertised catalog, never leak that raw XML/JSON as ordinary
    # content. OpenClaude has its own raw/XML fallback parser and could otherwise
    # execute an unadvertised hallucinated function behind this server's back.
    if effective_tools and not tool_calls and has_complete_tool_call(text):
        raise HTTPException(
            status_code=502,
            detail=(
                "Model produced a complete but invalid or unadvertised tool call; "
                "refusing to expose it as plain text to the tool executor."
            ),
        )

    message: dict[str, Any] = {"role": "assistant", "content": content or None}
    
    if tool_mode in {"required", "forced"} and not tool_calls:
        detail = (
            "Model failed to produce a valid required tool call. "
            "No plain-text success response was returned because OpenClaude "
            "requested tool execution."
        )
        if completion_tokens >= bounded_max_tokens:
            detail += " Generation reached the output-token limit."
        raise HTTPException(status_code=502, detail=detail)

    finish_reason = "stop"
    if tool_calls:
        message["tool_calls"] = tool_calls
        finish_reason = "tool_calls"
    elif completion_tokens >= bounded_max_tokens:
        finish_reason = "length"

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": MODEL,
        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "model": MODEL,
        "model_loaded": model is not None,
        "context_length": MAX_CONTEXT_TOKENS,
        "yarn_enabled": YARN_ENABLED,
        "yarn_factor": YARN_FACTOR if YARN_ENABLED else 1.0,
        "zero_gpu_size": ZERO_GPU_SIZE,
    }


def models() -> dict[str, Any]:
    return {
        "object": "list",
        "data": [
            {
                "id": model_id,
                "object": "model",
                "owned_by": "Erinaldorodrigues",
                "context_length": MAX_CONTEXT_TOKENS,
                "max_input_tokens": MAX_CONTEXT_TOKENS,
                "max_output_tokens": MAX_NEW_TOKENS,
            }
            for model_id in MODEL_ALIASES
        ],
    }


def chat_completions(request: ChatCompletionRequest):
    completion = _completion_payload(request)
    if not request.stream:
        return JSONResponse(content=completion)

    choice = completion["choices"][0]
    chunk_id = completion["id"]

    def events():
        first = {
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": completion["created"],
            "model": MODEL,
            "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
        }
        yield f"data: {json.dumps(first)}\n\n"
        
        delta: dict[str, Any] = {}
        if choice["message"].get("content"):
            delta["content"] = choice["message"]["content"]
        if choice["message"].get("tool_calls"):
            delta["tool_calls"] = indexed_tool_calls(
                choice["message"]["tool_calls"]
            )
            
        body = {**first, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}
        yield f"data: {json.dumps(body)}\n\n"
        
        final = {**first, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}
        yield f"data: {json.dumps(final)}\n\n"

        if request.stream_options and request.stream_options.get("include_usage") is True:
            usage_chunk = {
                "id": chunk_id,
                "object": "chat.completion.chunk",
                "created": completion["created"],
                "model": MODEL,
                "choices": [],
                "usage": completion["usage"],
            }
            yield f"data: {json.dumps(usage_chunk)}\n\n"

        yield "data: [DONE]\n\n"

    return StreamingResponse(
        events(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


async def _chat_completions_with_request_context(
    http_request: Request,
    parsed_request: ChatCompletionRequest,
):
    """Preserve the incoming HF/Gradio request through the thread boundary.

    ZeroGPU attributes quota through ``gradio.context.LocalContext.request``.
    Custom OpenAI routes bypass Gradio's normal event-listener setup, so bind
    the Starlette request explicitly before entering ``run_in_threadpool``.
    AnyIO/Starlette propagate contextvars into the worker thread, allowing the
    ``@spaces.GPU`` wrapper to see the HF proxy's ``x-ip-token`` header.
    """
    context_token = LocalContext.request.set(http_request)
    try:
        return await run_in_threadpool(chat_completions, parsed_request)
    finally:
        LocalContext.request.reset(context_token)


def _zerogpu_limit_response(error: Exception) -> JSONResponse | None:
    """Turn ZeroGPU quota/capacity rejections into an actionable 429."""
    message = str(error) or error.__class__.__name__
    lowered = message.casefold()
    quota_markers = (
        "space app has reached its gpu limit",
        "zerogpu quota exceeded",
        "gpu quota",
        "out of quota",
    )
    if not any(marker in lowered for marker in quota_markers):
        return None
    return JSONResponse(
        status_code=429,
        content={
            "error": {
                "message": (
                    "ZeroGPU rejected the GPU request because quota/capacity is "
                    "unavailable. For direct hf.space API calls, send a valid "
                    "Hugging Face token as Authorization: Bearer hf_... so the "
                    "call is charged to the caller's ZeroGPU quota. This Space "
                    "uses xlarge, which consumes quota at 2x. Original scheduler "
                    f"message: {message}"
                )
            }
        },
    )


demo = gr.Interface(
    fn=gerar,
    inputs=[
        gr.Textbox(label="Messages JSON"),
        gr.Number(value=DEFAULT_TEMPERATURE, label="Temperature"),
        gr.Number(value=512, label="Max Tokens"),
        gr.Textbox(value="[]", label="Tools JSON"),
        gr.Checkbox(value=True, label="Stop after first complete tool call"),
    ],
    outputs="text",
    title="Qwen2.5-Coder-32B AWQ OpenAI-compatible ZeroGPU Backend",
)

class OpenAIRouteMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        path = request.url.path.rstrip("/") or "/"
        
        if path == "/health" and request.method == "GET":
            return JSONResponse(health())
            
        if path == "/web-search" and request.method == "GET":
            query = request.query_params.get("q", "").strip()
            if not query or len(query) > 500:
                return JSONResponse(status_code=400, content={"error": "invalid query"})
            try:
                return JSONResponse(await run_in_threadpool(search_web, query))
            except SearchUnavailable as error:
                return JSONResponse(
                    status_code=503,
                    content={"error": {"message": str(error) or "search unavailable"}},
                )
            except Exception as error:
                traceback.print_exc()
                return JSONResponse(
                    status_code=500,
                    content={"error": {"message": f"search error: {error}"}},
                )
                
        if path == "/v1/models" and request.method == "GET":
            return JSONResponse(models())
            
        if path == "/v1/chat/completions" and request.method == "POST":
            try:
                raw_request = await request.json()
                parsed_request = ChatCompletionRequest(**raw_request)
            except (json.JSONDecodeError, ValidationError, TypeError) as error:
                return JSONResponse(status_code=400, content={"error": {"message": str(error)}})
            try:
                return await _chat_completions_with_request_context(
                    request, parsed_request
                )
            except HTTPException as error:
                return JSONResponse(status_code=error.status_code, content={"error": {"message": error.detail}})
            except Exception as error:
                quota_response = _zerogpu_limit_response(error)
                if quota_response is not None:
                    return quota_response
                traceback.print_exc()
                return JSONResponse(
                    status_code=500,
                    content={"error": {"message": f"internal Space error: {str(error)}"}}
                )
                
        return await call_next(request)


import gradio.routes as _groutes
_original_create_app = _groutes.App.create_app

def _create_app_with_openai_routes(*args, **kwargs):
    created = _original_create_app(*args, **kwargs)
    created.add_middleware(OpenAIRouteMiddleware)
    return created

_groutes.App.create_app = staticmethod(_create_app_with_openai_routes)

demo.queue(default_concurrency_limit=1, max_size=8).launch(show_error=True, ssr_mode=False)