import os
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import spaces
import gradio as gr

# Official Public Nous Research Repository (Kabhi delete ya 401 nahi hoga)
MODEL_ID = "NousResearch/Hermes-3-Llama-3.1-8B"

print("1. Loading Tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

model = None

# 4-bit Quantization configuration for ZeroGPU (~5.5GB VRAM usage)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

def get_or_load_model():
    global model
    if model is None:
        print("2. Mounting 4-bit Model inside ZeroGPU context...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            quantization_config=bnb_config,
            device_map="auto",
            low_cpu_mem_usage=True
        )
    return model

# Pydantic Schemas for OpenAI compatibility
class Message(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = "Hermes-3-8B"
    messages: List[Message]
    max_tokens: Optional[int] = 2048
    temperature: Optional[float] = 0.6

# ZeroGPU Generation Function
@spaces.GPU(duration=120)
def generate_hermes(prompt_text: str, max_tokens: int, temperature: float) -> str:
    active_model = get_or_load_model()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = active_model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True if temperature > 0.0 else False,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    
    generated_tokens = outputs[0][inputs.input_ids.shape[-1]:]
    
    # skip_special_tokens=False zaroori hai taaki <tool_call> tags strip na hon
    return tokenizer.decode(generated_tokens, skip_special_tokens=False)

# UI element space ko active rakhne ke liye
with gr.Blocks(title="Hermes-3 8B Server") as demo:
    gr.Markdown("# 🚀 Hermes-3 8B ZeroGPU API Server")
    gr.Markdown("Active Endpoints:\n- `GET /health`\n- `POST /v1/chat/completions`")

app = demo.app

@app.get("/health")
def health_check():
    return {"status": "ok", "ready": True}

@app.post("/v1/chat/completions")
def chat_completions(req: ChatCompletionRequest):
    try:
        formatted_messages = [{"role": m.role, "content": m.content} for m in req.messages]
        
        # Hermes ChatML formatting
        prompt_text = tokenizer.apply_chat_template(
            formatted_messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        raw_output = generate_hermes(
            prompt_text=prompt_text,
            max_tokens=req.max_tokens,
            temperature=req.temperature
        )

        # Hermes stop tags trim karein
        clean_output = raw_output.replace("<|im_end|>", "").strip()

        return {
            "id": "chatcmpl-hermes",
            "object": "chat.completion",
            "model": req.model,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": clean_output
                    },
                    "finish_reason": "stop"
                }
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)