import os
import json
import torch
import torch.nn as nn
import sentencepiece as spm
import gradio as gr
from huggingface_hub import hf_hub_download
import spaces

# --- 1. CONFIGURATION & MODEL DOWNLOAD ---
REPO_ID = "yisberh/grade8-math-custom-transformer"

# Fetch preloaded cache files instantly from local disk memory
config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json", local_files_only=True)
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt", local_files_only=True)
spm_path = hf_hub_download(repo_id=REPO_ID, filename="unit1_test.model", local_files_only=True)

with open(config_path, "r", encoding="utf-8") as f:
    config = json.load(f)

sp = spm.SentencePieceProcessor(model_file=spm_path)

# --- 2. MODEL DEFINITION (Matching your exact architecture) ---
class Model(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.d = 256  # 256 embedding dimensions
        self.emb = nn.Embedding(vocab, self.d)
        self.pos = nn.Embedding(512, self.d)
        layer = nn.TransformerEncoderLayer(
            d_model=self.d,
            nhead=8,      # 8 attention heads
            dim_feedforward=1024, # 1024 FFN dimension
            batch_first=True,
            dropout=0.0   # 0.0 dropout for speed during inference
        )
        self.tr = nn.TransformerEncoder(layer, 6) # 6 Layers deep
        self.head = nn.Linear(self.d, vocab)

    def forward(self, x):
        n = x.size(1)
        pos = torch.arange(n, device=x.device)
        x = self.emb(x) + self.pos(pos)
        mask = torch.triu(
            torch.full((n, n), float("-inf"), device=x.device),
            diagonal=1
        )
        return self.head(self.tr(x, mask=mask))

# Initialize model container framework on CPU first
model = Model(sp.vocab_size())

# Strip any potential DataParallel state_dict prefix formatting cleanly
state_dict = torch.load(model_path, map_location="cpu")
if any(k.startswith("module.") for k in state_dict.keys()):
    state_dict = {k.replace("module.", ""): v for k in state_dict.items()}

model.load_state_dict(state_dict)
model.eval()

# --- 3. INFERENCE LOOP ROUTED VIA ZEROGPU ---
@spaces.GPU(duration=60)  # Mandated to prevent "No @spaces.GPU function detected"
@torch.inference_mode()
def generate_response(question):
    # Protect against early app validation/mock checks passing empty inputs
    if not question or not isinstance(question, str) or not question.strip():
        return "Please ask a question about Grade 8 Mathematics Unit 1."

    # Identify device assigned dynamically by the Hugging Face manager
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)

    # Encode instructions using your native data pipeline format
    ids = [sp.bos_id()] + sp.encode_as_ids(f"[INST] {question.strip()} [/INST]")
    x = torch.tensor([ids], device=device)
    max_new_tokens = 256 # Higher generation limits to support extended answers

    for _ in range(max_new_tokens):
        if x.size(1) >= 512: # Strict constraint boundary matching sequence history size
            break
            
        logits = model(x)
        nxt = logits[:, -1].argmax(-1).item()
        
        nxt_tensor = torch.tensor([[nxt]], device=device)
        x = torch.cat([x, nxt_tensor], dim=1)
        
        if nxt == sp.eos_id():
            break

    text = sp.decode(x[0].tolist())
    
    # Split using the explicit token tags matching your SFT layout formatting
    if "[/INST]" in text:
        return text.split("[/INST]")[-1].strip()
    return text

# --- 4. GRADIO INTERFACE LAYOUT CONFIGURATION ---
demo = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Enter your Grade 8 Mathematics question here... (Supports English & Amharic)",
        label="Math Question"
    ),
    outputs=gr.Textbox(label="Model Response"),
    title="Grade 8 Mathematics AI Assistant (5M Params)",
    description="Custom transformer model tracking Grade 8 Math Unit 1 roadmap objectives."
)

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