try:
    import spaces
except ImportError:
    spaces = None

import gradio as gr
import torch
import qtensor
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

device = "cuda" if torch.cuda.is_available() else "cpu"
use_triton = torch.cuda.is_available()

# 1. Load Instruct model architecture and tokenizer
base_model_id = "unsloth/Llama-3.2-1B-Instruct"
healed_repo_id = "trentzap/Llama-3.2-1B-QTensor-FP8"

print(f"[QTensor Space] Loading base architecture {base_model_id} on {device}...")
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

raw_model = AutoModelForCausalLM.from_pretrained(
    base_model_id, 
    dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, 
    device_map=device
)

# 2. Compress MLP layers (66.7% of parameters) into QTensor FP8 MPOLinear layers
print("[QTensor Space] Structuring QTensor FP8 MPOLinear layers...")
model = qtensor.compress(raw_model, precision="fp8", chi=512, use_triton=use_triton, target_layers=["gate_proj", "up_proj", "down_proj"])

# 3. Load pre-healed weights from Hugging Face Hub
try:
    print(f"[QTensor Space] Downloading healed weights from {healed_repo_id}...")
    weights_path = hf_hub_download(repo_id=healed_repo_id, filename="model.safetensors")
    state_dict = load_file(weights_path)
    model.load_state_dict(state_dict, strict=False)
    print("[QTensor Space] Successfully loaded pre-healed QTensor weights!")
except Exception as e:
    print(f"[QTensor Space] Notice: Running with default weights ({e})")

model.eval()

# 4. Define chat response generator with chat template formatting
def _generate(prompt, history):
    messages = []
    if history:
        for user_msg, assistant_msg in history:
            messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": prompt})
    
    formatted_prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return response

if spaces is not None:
    _generate_gpu = spaces.GPU(_generate)
else:
    _generate_gpu = _generate

def generate_response(prompt, history):
    return _generate_gpu(prompt, history)

# 5. Launch Gradio Chat Interface
demo = gr.ChatInterface(
    fn=generate_response,
    title="QTensor LLaMA-3.2 1B Chat (FP8 / 0.68GB VRAM)",
    description="Running on QTensor 160-bit MPO + FP8 Stacked Quantization and Triton SRAM Fusion kernels!"
)

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