import os
import gradio as gr
import torch
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from peft import PeftModel
from huggingface_hub import hf_hub_download

# --- 1. Configuration ---
# Ensure 'HF_TOKEN' is set in your Space's Settings > Secrets
HF_TOKEN = os.getenv("HF_TOKEN")

# Model IDs
LARGE_BASE = "meta-llama/Llama-3.1-8B" 
LARGE_ADAPTER = "tuochao/Llama-3.1-8B-Proactive-Big-Peft"

SMALL_BASE = "meta-llama/Llama-3.2-1B"
# YOUR NEW REPO (Contains adapter + classifier_weights.pth)
SMALL_ADAPTER = "xubayer/Llama-3.2-1B-Proactive-Small-Complete"

# 4-bit Quantization Config (Crucial for fitting both models on T4 GPU)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
)

# --- 2. Load Small Classifier (Custom Architecture) ---
print(f"Loading Classifier Base & Adapter: {SMALL_ADAPTER}...")

# A. Load Base Model (CausalLM) + LoRA Adapter
# We load as CausalLM because the PEFT adapter was trained on a CausalLM base.
cl_base = AutoModelForCausalLM.from_pretrained(
    SMALL_BASE,
    quantization_config=bnb_config,
    device_map="auto",
    token=HF_TOKEN
)
classifier_model = PeftModel.from_pretrained(cl_base, SMALL_ADAPTER, token=HF_TOKEN)
classifier_tokenizer = AutoTokenizer.from_pretrained(SMALL_BASE, token=HF_TOKEN)

# Fix for Llama 3 tokenizer
if classifier_tokenizer.pad_token is None:
    classifier_tokenizer.pad_token = classifier_tokenizer.eos_token
    classifier_model.config.pad_token_id = classifier_model.config.eos_token_id

# B. Reconstruct Custom Classifier Head
# This manually adds the Linear layer (hidden_size -> 2) that Unsloth used.
hidden_size = classifier_model.config.hidden_size
classifier_head = torch.nn.Linear(hidden_size, 2).to(classifier_model.device)

# C. Download and Load Custom Weights
try:
    print("Downloading custom classifier head weights...")
    weights_path = hf_hub_download(
        repo_id=SMALL_ADAPTER, 
        filename="classifier_weights.pth", 
        token=HF_TOKEN
    )
    # Load weights
    state_dict = torch.load(weights_path, map_location=classifier_model.device)
    classifier_head.load_state_dict(state_dict)
    print("✅ Custom classifier weights loaded successfully.")
except Exception as e:
    print(f"⚠️ CRITICAL ERROR: Could not load 'classifier_weights.pth'.\nDetails: {e}")
    # We don't crash here, but inference will be random if this fails.

# --- 3. Load Large Generator ---
print(f"Loading Generator: {LARGE_ADAPTER}...")
gen_base = AutoModelForCausalLM.from_pretrained(
    LARGE_BASE,
    quantization_config=bnb_config,
    device_map="auto",
    token=HF_TOKEN
)
generator_model = PeftModel.from_pretrained(gen_base, LARGE_ADAPTER, token=HF_TOKEN)
generator_tokenizer = AutoTokenizer.from_pretrained(LARGE_BASE, token=HF_TOKEN)

if generator_tokenizer.pad_token is None:
    generator_tokenizer.pad_token = generator_tokenizer.eos_token

# --- 4. Inference Logic ---
def predict(history_text):
    """
    Pipeline:
    1. Append |SILENCE > marker.
    2. Extract hidden state from 1B model.
    3. Pass through custom classifier head.
    4. If 'Whisper' (Class 1) > Threshold, run 8B generator.
    """
    
    # Input formatting
    classifier_input = history_text.strip()
    # Assuming your model was trained to see this marker at the end
    if not classifier_input.endswith("|SILENCE >"):
         classifier_input += " |SILENCE >"

    # --- Stage 1: Classifier ---
    inputs = classifier_tokenizer(
        classifier_input, 
        return_tensors="pt",
        truncation=True,
        max_length=2048
    ).to(classifier_model.device)

    with torch.no_grad():
        # Run base model to get hidden states
        outputs = classifier_model(
            **inputs, 
            output_hidden_states=True
        )
        # Get the last hidden state of the last token (the '>' in |SILENCE >)
        last_hidden_state = outputs.hidden_states[-1]
        last_token_state = last_hidden_state[:, -1, :]
        
        # Pass through our custom linear head
        logits = classifier_head(last_token_state)
        
        # Convert to probability
        probs = torch.nn.functional.softmax(logits, dim=-1)
        whisper_score = probs[0][1].item() # Probability of Class 1 (Whisper)

    status_msg = f"Confidence: {whisper_score:.2%} (Threshold: 50%)"
    
    # --- Stage 2: Generator ---
    whisper_content = ""
    threshold = 0.5 
    
    if whisper_score > threshold:
        status_msg += " -> 🟢 TRIGGERED"
        
        # Generator sees the same history
        gen_inputs = generator_tokenizer(history_text, return_tensors="pt").to(generator_model.device)
        
        with torch.no_grad():
            gen_out = generator_model.generate(
                **gen_inputs,
                max_new_tokens=40,   # Keep response short
                do_sample=True,      # Creative generation
                temperature=0.6,
                top_p=0.9,
                pad_token_id=generator_tokenizer.pad_token_id
            )
        
        full_text = generator_tokenizer.decode(gen_out[0], skip_special_tokens=True)
        
        # Clean up: Remove the original conversation history from the output
        # (This logic assumes the model repeats the prompt, which CausalLMs do)
        if full_text.startswith(history_text):
             whisper_content = full_text[len(history_text):].strip()
        else:
             whisper_content = full_text # Fallback
             
    else:
        status_msg += " -> 🔴 SILENT"
        whisper_content = "(No whisper needed)"

    return status_msg, whisper_content

# --- 5. UI Setup ---
with gr.Blocks(title="LlamaPIE Final Demo") as demo:
    gr.Markdown("# 🥧 LlamaPIE: Proactive In-Ear Assistant")
    gr.Markdown("A dual-model pipeline: **1B Classifier** detects silence, **8B Generator** whispers suggestions.")
    
    with gr.Row():
        input_box = gr.Textbox(
            label="Conversation Context", 
            lines=4, 
            placeholder="User: I'm bored.\nAI: Do you want to watch a movie? |SILENCE >",
            value="User: I forgot to buy milk.\nAI: That's annoying.\nUser: Yeah, I guess I'll have black coffee. |SILENCE >"
        )
    
    btn = gr.Button("⚡ Run Inference Pipeline", variant="primary")
    
    with gr.Row():
        d_box = gr.Textbox(label="Stage 1: Decision (Confidence)")
        w_box = gr.Textbox(label="Stage 2: Generated Whisper")

    btn.click(fn=predict, inputs=input_box, outputs=[d_box, w_box])

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