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

from __future__ import annotations

import json
import os
import time
import uuid
import traceback
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
import spaces
import torch
from fastapi import HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, 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,
)
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"),
)

MAX_CONTEXT_TOKENS = int(os.getenv("MAX_CONTEXT_TOKENS", "16384"))
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "2048"))
MAX_TOOL_CALL_TOKENS = int(os.getenv("MAX_TOOL_CALL_TOKENS", "2048"))
MAX_TEMPERATURE = float(os.getenv("MAX_TEMPERATURE", "0.2"))
PRESERVED_PREFIX_TOKENS = int(os.getenv("PRESERVED_PREFIX_TOKENS", "4096"))

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


def _ensure_model_loaded() -> None:
    """Carrega o modelo uma única vez, já dentro do contexto com GPU real."""
    global model
    if model is not None:
        return
    print(f"Loading {MODEL} on ZeroGPU (NATIVE AWQ)...", flush=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL,
        torch_dtype="auto",
        device_map="auto",
        low_cpu_mem_usage=True,
    )
    model.eval()
    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,
    )
    return duration + COLD_START_BUFFER_SECONDS


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)


class StopAfterToolCall(StoppingCriteria):
    def __init__(self, prompt_length: int) -> None:
        self.prompt_length = prompt_length

    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))
        return torch.tensor(completed, dtype=torch.bool, device=input_ids.device)


@spaces.GPU(duration=_gpu_duration)
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)
    tool_mode = _tool_protocol_active(messages) or bool(tools)
    template_kwargs: dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": True,
    }
    if tools:
        template_kwargs["tools"] = tools

    try:
        prompt = tokenizer.apply_chat_template(messages, **template_kwargs)
    except Exception as template_error:
        print(f"Jinja Template Warning: {template_error}. Applying fallback.", flush=True)
        template_kwargs.pop("tools", None)
        prompt = tokenizer.apply_chat_template(messages, **template_kwargs)

    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:
        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 or 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])]
        )

    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]]
    temperature: float = 0.2
    max_tokens: int | None = None
    max_completion_tokens: int | None = None
    stream: bool = False
    tools: list[dict[str, Any]] | None = None
    tool_choice: Any = None
    parallel_tool_calls: bool | None = None


def _completion_payload(request: ChatCompletionRequest) -> dict[str, Any]:
    if request.model not in {
        MODEL,
        "qwen-coder",
        "qwen3-coder",
        "qwen2.5-coder-32b",
        "qwen2.5-coder-14b",
    }:
        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 [])
    state_controls_choice = request.tool_choice is None or (
        isinstance(request.tool_choice, str)
        and request.tool_choice.casefold() == "auto"
    )
    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)
                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
                else None
            ),
        )
        if instruction
    ]
    instruction = "\n\n".join(instructions) if instructions else None
    max_tokens = request.max_completion_tokens or request.max_tokens or MAX_NEW_TOKENS
    
    if effective_tools:
        max_tokens = min(max_tokens, MAX_TOOL_CALL_TOKENS)
    temperature = min(max(float(request.temperature), 0.01), MAX_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
    
    text = gerar(
        json.dumps(prompt_messages),
        temperature,
        _bounded_output_tokens(max_tokens),
        json.dumps(effective_tools, ensure_ascii=False),
        request.parallel_tool_calls is not True,
    )
    
    if effective_tools:
        tool_calls, content = extract_tool_calls(text, tool_names(effective_tools))
        if request.parallel_tool_calls is False:
            tool_calls = tool_calls[:1]
    else:
        tool_calls, content = [], text
        
    message: dict[str, Any] = {"role": "assistant", "content": content or None}
    
    finish_reason = "stop"
    if tool_calls:
        message["tool_calls"] = tool_calls
        finish_reason = "tool_calls"
    elif effective_tools and has_complete_tool_call(text):
        finish_reason = "stop"
    elif tool_mode in {"required", "forced"}:
        finish_reason = "stop"

    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": 0, "completion_tokens": 0, "total_tokens": 0},
    }


def health() -> dict[str, str]:
    return {"status": "ok", "model": MODEL}


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 dict.fromkeys(("qwen2.5-coder-32b", MODEL))
        ],
    }


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"
        yield "data: [DONE]\n\n"

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


demo = gr.Interface(
    fn=gerar,
    inputs=[
        gr.Textbox(label="Messages JSON"),
        gr.Number(value=0.2, label="Temperature"),
        gr.Number(value=512, label="Max Tokens"),
    ],
    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 Exception:
                return JSONResponse(status_code=500, content={"error": "search 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 chat_completions(parsed_request)
            except HTTPException as error:
                return JSONResponse(status_code=error.status_code, content={"error": {"message": error.detail}})
            except Exception as error:
                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)