import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. Download the highly optimized Llama 3.2 3B model from Hugging Face repository
print("Downloading model...")
model_path = hf_hub_download(
    repo_id="bartowski/Llama-3.2-3B-Instruct-GGUF",
    filename="Llama-3.2-3B-Instruct-Q4_K_M.gguf"
)

# 2. Initialize the model on the CPU
print("Initializing model...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,       # Context length
    n_threads=2       # Utilizes both free CPU cores fully
)

# 3. Define the chatbot logic
def respond(message, chat_history):
    # Format the prompt to match Llama 3.2 structural rules
    formatted_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
    formatted_prompt += "You are a helpful, direct, and honest AI assistant.<|eot_id|>"
    
    # Inject chat history so the bot remembers the conversation context
    for turn in chat_history:
        # Check if history is structured as objects or dicts (Gradio 6 style)
        user_msg = turn.get("text") if isinstance(turn, dict) else turn
        bot_msg = turn.get("text") if isinstance(turn, dict) else turn
        
        if user_msg:
            formatted_prompt += f"<|start_header_id|>user<|end_header_id|>\n{user_msg}<|eot_id|>"
        if bot_msg:
            formatted_prompt += f"<|start_header_id|>assistant<|end_header_id|>\n{bot_msg}<|eot_id|>"
            
    # Add the newest user message
    formatted_prompt += f"<|start_header_id|>user<|end_header_id|>\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
    
    # Generate response tokens streamingly
    output = llm(
        formatted_prompt,
        max_tokens=512,
        stop=["<|eot_id|>"],
        stream=True
    )
    
    token_accumulator = ""
    for token in output:
        token_text = token["choices"]["text"]
        token_accumulator += token_text
        yield token_accumulator

# 4. Create the web dashboard layout using a global theme block wrapper
with gr.Blocks(theme="soft") as demo:
    gr.ChatInterface(
        fn=respond,
        title="🤖 Free Llama 3.2 CPU Chatbot",
        description="Running 24/7/365 for free on Hugging Face Spaces using CPU inference.",
        examples=["Explain quantum computing simply.", "Write a short poem about coding."]
    )

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