"""Hugging Face Space — fine-tuned Flutter coder (ZeroGPU, free tier).

Deploy: create a new Gradio Space, upload this folder, set hardware to ZeroGPU.
Set HF Space URL in the Flutter IDE Settings (HUGGINGFACE_SPACE_URL).
"""

import os

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MODEL_ID = os.environ.get("HF_MODEL_ID", "malek391/my-custom-flutter-coder")
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_ACCESS_TOKEN")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "2048"))

_tokenizer = None
_model = None


def load_model():
    """Load on first GPU request — ZeroGPU has no CUDA at container startup."""
    global _tokenizer, _model
    if _model is not None:
        return _tokenizer, _model

    _tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        trust_remote_code=True,
    )
    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
    )
    _model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        trust_remote_code=True,
        quantization_config=quant_config,
        device_map="cuda",
    )
    return _tokenizer, _model


def gpu_duration(system_prompt: str, user_prompt: str, temperature: float, *args, **kwargs) -> int:
    """Keep declared duration under ZeroGPU free-tier per-call cap (~180s effective)."""
    prompt_len = len(system_prompt or "") + len(user_prompt or "")
    # Shorter prompts → lower reservation (better queue priority, less quota blocked)
    return min(90, 60 + prompt_len // 800)


@spaces.GPU(duration=gpu_duration)
def chat(system_prompt: str, user_prompt: str, temperature: float) -> str:
    """OpenAI-style chat via system + user strings (Gradio API name: chat)."""
    tokenizer, model = load_model()

    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    messages.append({"role": "user", "content": user_prompt.strip()})

    if hasattr(tokenizer, "apply_chat_template"):
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    else:
        prompt = f"{system_prompt}\n\n{user_prompt}"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    temp = max(0.05, min(float(temperature), 1.0))

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            temperature=temp,
            do_sample=temp > 0.05,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = output[0][inputs.input_ids.shape[1] :]
    return tokenizer.decode(generated, skip_special_tokens=True)


demo = gr.Interface(
    fn=chat,
    inputs=[
        gr.Textbox(label="System prompt", lines=4, placeholder="System instructions…"),
        gr.Textbox(label="User prompt", lines=12, placeholder="User message…"),
        gr.Slider(0.05, 1.0, value=0.25, step=0.05, label="Temperature"),
    ],
    outputs=gr.Textbox(label="Model response", lines=20),
    title="My Custom Flutter Coder",
    description=(
        f"Fine-tuned Qwen2.5-Coder served from `{MODEL_ID}` on ZeroGPU (4-bit). "
        "First request may take ~1 min while the model loads."
    ),
    api_name="chat",
)

if __name__ == "__main__":
    demo.launch()