# app.py
import os
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces  # Mandatory library for Hugging Face ZeroGPU

MODEL_ID = "DevStudio-AI/Devstudio-Coder-1.5B"
HF_TOKEN = os.environ.get("HF_TOKEN")

print("Loading tokenizer and base model...")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen2.5-Coder-1.5B-Instruct"
)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    subfolder="models/final_merged",
    torch_dtype=torch.float16,
    device_map="cpu",
    token=HF_TOKEN
)

print("Model successfully loaded on CPU. Awaiting ZeroGPU allocation...")


@spaces.GPU
def generate_code(prompt, temperature, max_tokens):
    try:
        # Move model to GPU inside the ZeroGPU context
        model.to("cuda")

        system_prompt = (
            "You are DevStudio-1.5B, an in-editor coding assistant developed by DevStudio AI. "
            "You are a highly specialized master of modern single-file HTML and Tailwind CSS designs."
        )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]

        # Apply ChatML formatting
        formatted_prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = tokenizer(
            formatted_prompt,
            return_tensors="pt"
        ).to("cuda")

        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id
        )

        # Decode only the generated part
        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        result = tokenizer.decode(
            generated_ids,
            skip_special_tokens=True
        )

        # ---------------------------
        # Post-process escaped output
        # ---------------------------
        result = result.replace("\\n", "\n")
        result = result.replace("\\t", "\t")
        result = result.replace('\\"', '"')
        result = result.replace("\\'", "'")

        # Remove escaped markdown fences if present
        result = result.replace("\\`", "`")

        return result

    except Exception as e:
        return f"Error during generation: {str(e)}"

    finally:
        # Free GPU memory after each request
        if torch.cuda.is_available():
            model.to("cpu")
            torch.cuda.empty_cache()


demo = gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(
            label="Prompt",
            placeholder="Enter your prompt here...",
            lines=8
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.3,
            label="Temperature"
        ),
        gr.Slider(
            minimum=64,
            maximum=2048,
            value=1024,
            step=64,
            label="Max Tokens"
        )
    ],
    outputs=gr.Textbox(
        label="Generated Code",
        lines=30
    ),
    title="DevStudio-1.5B API Engine",
    description="Static HTML + Tailwind CSS specialized completion endpoint running on ZeroGPU."
)

demo.launch()