"""
Apex-64M: Interactive Chat Assistant Web Application
Deployable directly on Hugging Face Spaces (CPU or GPU)
"""

import os
import torch
import gradio as gr
from config import GPTConfig
from model import SmallGPT
from tokenizer import get_tokenizer

# 1. Device & Tokenizer Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = get_tokenizer()
eot_id = tokenizer.eot_token  # 50256

# 2. Model Initialization & Weight Loading
config = GPTConfig()
model = SmallGPT(config).to(device)

# Look for SFT chat weights first, then base weights
CHECKPOINT_CANDIDATES = [
    "sft_model.pt",
    "slm_chat_assistant.pt",
    "best_model.pt",
    "slm_weights.pt",
    "model.pt"
]

loaded_ckpt = None
for ckpt in CHECKPOINT_CANDIDATES:
    if os.path.exists(ckpt):
        print(f"[*] Loading model weights from: {ckpt}")
        state = torch.load(ckpt, map_location=device)
        if isinstance(state, dict) and "model_state_dict" in state:
            state = state["model_state_dict"]
        model.load_state_dict(state)
        loaded_ckpt = ckpt
        break

if loaded_ckpt:
    print(f"[✓] Successfully loaded {loaded_ckpt} on {device}")
else:
    print("[!] Warning: No checkpoint found. Running with randomly initialized weights.")

model.eval()

# 3. Multi-Turn Inference Logic
def generate_chat_response(
    user_input: str,
    chat_history: list,
    temperature: float = 0.7,
    top_p: float = 0.9,
    repetition_penalty: float = 1.2
) -> tuple[str, list]:
    if not user_input or not user_input.strip():
        return "", chat_history or []
        
    chat_history = chat_history or []
    
    # Build prompt from multi-turn history
    history_prompts = []
    for u, a in chat_history:
        history_prompts.append(f"### Instruction:\n{u}\n\n### Response:\n{a}")
        
    history_prompts.append(f"### Instruction:\n{user_input.strip()}\n\n### Response:\n")
    full_prompt = "\n\n".join(history_prompts)
    
    # Sliding context window (keep last 380 tokens)
    tokens = tokenizer.encode(full_prompt)
    if len(tokens) > 380:
        tokens = tokens[-380:]
        full_prompt = tokenizer.decode(tokens)
        
    idx = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
    
    with torch.no_grad():
        for _ in range(130):
            idx_cond = idx if idx.size(1) <= config.context_length else idx[:, -config.context_length:]
            logits, _ = model(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            
            # Repetition Penalty
            for token_id in set(idx[0].tolist()):
                if logits[0, token_id] > 0:
                    logits[0, token_id] /= repetition_penalty
                else:
                    logits[0, token_id] *= repetition_penalty
            
            # Top-P Nucleus Filtering
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0
            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            logits[:, indices_to_remove] = -float("Inf")
            
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            # Stop immediately on <|endoftext|>
            if next_token.item() == eot_id:
                break
                
            idx = torch.cat((idx, next_token), dim=1)
            
            # Early stop if model predicts new prompt header
            current_text = tokenizer.decode(idx[0].tolist())
            if "### Instruction:" in current_text.split("### Response:\n")[-1]:
                break
                
    full_text = tokenizer.decode(idx[0].tolist())
    
    # Extract clean answer
    answer = full_text.split("### Response:\n")[-1]
    answer = answer.split("### Instruction:")[0]
    answer = answer.split("### Context:")[0]
    answer = answer.replace("<|endoftext|>", "").strip()
    
    chat_history.append((user_input, answer))
    return "", chat_history


def clear_history():
    return "", []


# 4. Gradio UI Interface
with gr.Blocks(title="Apex-64M Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🤖 Apex-64M Assistant
        **A 63.8M Parameter Small Language Model (SLM) trained 100% from scratch on ~1.28B tokens of C4 English + Fine-Tuned on Dolly 15k.**
        
        *Zero pre-trained weights used. Built entirely in PyTorch with custom FlashAttention.*
        """
    )
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=460, label="Conversation")
            msg = gr.Textbox(
                placeholder="Ask Apex anything (e.g. What is the difference between a stack and a queue?)...",
                label="Your Message",
                lines=1
            )
            with gr.Row():
                send_btn = gr.Button("Send 🚀", variant="primary")
                clear_btn = gr.Button("Clear Memory 🧹")
                
            gr.Examples(
                examples=[
                    "What is the difference between a stack and a queue?",
                    "Give me 3 tips for effective time management.",
                    "Explain how neural networks work in simple terms.",
                    "Write a short story about an astronaut exploring Mars."
                ],
                inputs=msg
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Generation Controls")
            temperature = gr.Slider(
                minimum=0.1,
                maximum=1.5,
                value=0.7,
                step=0.05,
                label="Temperature (Creativity)",
                info="Lower = focused & factual, Higher = creative"
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top-P (Nucleus Sampling)",
                info="Limits token selection to top cumulative probability"
            )
            repetition_penalty = gr.Slider(
                minimum=1.0,
                maximum=2.0,
                value=1.2,
                step=0.05,
                label="Repetition Penalty",
                info="Discourages repeating phrases"
            )
            
            gr.Markdown(
                """
                ---
                ### 📊 Model Architecture:
                - **Model Name:** Apex-64M
                - **Parameters:** ~63.8 Million
                - **Layers:** 12 Transformer Blocks
                - **Heads:** 8 Attention Heads (head_dim=64)
                - **Context Window:** 512 tokens
                - **Vocabulary:** 50,257 (tiktoken BPE)
                """
            )

    # Event Handlers
    send_btn.click(
        generate_chat_response,
        inputs=[msg, chatbot, temperature, top_p, repetition_penalty],
        outputs=[msg, chatbot]
    )
    msg.submit(
        generate_chat_response,
        inputs=[msg, chatbot, temperature, top_p, repetition_penalty],
        outputs=[msg, chatbot]
    )
    clear_btn.click(clear_history, outputs=[msg, chatbot])

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