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

# 1. DOWNLOAD MODEL
model_path = hf_hub_download(
    repo_id="NousResearch/Hermes-3-Llama-3.2-3B-GGUF", 
    filename="Hermes-3-Llama-3.2-3B.Q4_K_M.gguf", 
    local_dir="models"
)

# 2. LATENCY-OPTIMIZED ENGINE
llm = Llama(
    model_path=model_path,
    n_ctx=1024,      # Reduced to 1024 for faster "thought" time on CPU
    n_threads=2,     # Matches HF Free Tier cores
    n_batch=32,      # SMALLER batch = much faster first token appearance
    verbose=False
)

# 3. ROBUST CHAT LOGIC
def respond(message, history, system_prompt, temperature, max_tokens, top_p):
    # This logic works whether history is a list of lists OR a list of dicts
    messages = [{"role": "system", "content": system_prompt}]
    
    for interaction in history:
        if isinstance(interaction, dict):
            messages.append(interaction)
        elif isinstance(interaction, (list, tuple)) and len(interaction) == 2:
            messages.append({"role": "user", "content": interaction[0]})
            messages.append({"role": "assistant", "content": interaction[1]})
    
    messages.append({"role": "user", "content": message})

    # Call the model
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True
    )

    token_text = ""
    for chunk in response:
        if 'choices' in chunk and len(chunk['choices']) > 0:
            delta = chunk['choices'][0]['delta']
            if 'content' in delta:
                token_text += delta['content']
                yield token_text

# 4. COMPATIBLE GRADIO UI
demo = gr.ChatInterface(
    fn=respond,
    title="Instant Novelist AI",
    additional_inputs=[
        gr.Textbox(value="You are a gritty novelist. Write in a vivid style.", label="System Prompt"),
        gr.Slider(0.1, 2.0, 0.8, label="Temperature"),
        gr.Slider(128, 1024, 256, label="Max New Tokens"), # Lowered for snappiness
        gr.Slider(0.1, 1.0, 0.9, label="Top-p"),
    ]
)

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