import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

# ── CONFIGURATION ─────────────────────────────────────────────────────────────
MODEL_ID = "Havoc999/tiny-chatbot" 

print("Loading tokenizer and model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, 
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)
model.eval()

# ── INFERENCE ENGINE ──────────────────────────────────────────────────────────
def respond(message, history):
    """
    message: The current user prompt string.
    history: A list of dicts/lists representing the chat history.
    """
    # Format past conversation as context for the Alpaca template
    context_str = ""
    if len(history) > 0:
        context_str = "Past conversation history:\n"
        # Keep last 3 turns to avoid hitting max token limits
        for turn in history[-3:]: 
            # Gradio 6 history can be parsed safely via dict or index access
            user_msg = turn.get("user") if isinstance(turn, dict) else turn
            bot_msg = turn.get("options") if isinstance(turn, dict) else turn
            if user_msg and bot_msg:
                context_str += f"User: {user_msg}\nAssistant: {bot_msg}\n"
    
    # Build the Alpaca format string
    input_section = f"### Input:\n{context_str}\n\n" if context_str else ""
    prompt = (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        f"### Instruction:\n{message}\n\n"
        f"{input_section}"
        "### Response:\n"
    )

    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
    # TextIteratorStreamer yields tokens on the fly
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        
    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.15,
        pad_token_id=tokenizer.eos_token_id
    )

    # Run generation inside a background thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Stream the output chunks back to ChatInterface
    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message

# ── GRADIO INTERFACE ──────────────────────────────────────────────────────────
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 TinyLlama Chatbot")
    gr.Markdown("A ChatGPT-style interface for the fine-tuned `tiny-chatbot` model.")
    
    gr.ChatInterface(
        fn=respond,
        textbox=gr.Textbox(
            placeholder="Type a message...", 
            container=False, 
            scale=7
        )
    )

if __name__ == "__main__":
    # In Gradio 6.0+, theme configurations are passed strictly inside launch()
    demo.queue().launch(theme=gr.themes.Soft())