from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import os

def load_llm():
    print("--- Downloading/Loading Qwen2.5 1.5B GGUF (4-bit Quantized) ---")
    
    # This downloads the heavily compressed CPU version of Qwen
    # It caches it locally, so it only downloads once!
    model_path = hf_hub_download(
        repo_id="Qwen/Qwen2.5-1.5B-Instruct-GGUF",
        filename="qwen2.5-1.5b-instruct-q4_k_m.gguf"
    )
    
    # Load into standard CPU RAM
    # n_threads should match your physical CPU cores (e.g., 4 or 8)
    llm = Llama(
        model_path=model_path,
        n_ctx=2048,       # How much text it can remember at once
        n_threads=4,      # Number of CPU cores to use
        verbose=False     # Hides the messy C++ logs
    )
    print("--- CPU Model Loaded Successfully ---")
    return llm

def generate_answer(llm, retrieved_context, query):
    # Qwen-specific ChatML format for the CPU engine
    # This guarantees the model stays in 'Julia Assistant' mode
    prompt = (
        f"<|im_start|>system\n"
        f"You are the 'Ask Julia' assistant for diffusion.fyi. "
        f"Your tone is encouraging and academic. "
        f"You MUST use LaTeX for math (e.g., $x$ or $$y=x$$). "
        f"CRITICAL INSTRUCTION: You may ONLY answer based on the provided CONTEXT. "
        f"If the exact name or concept asked by the user is NOT in the CONTEXT, you MUST reply with EXACTLY: 'I don't have that specific information in my notes.' "
        f"DO NOT guess, DO NOT talk about other constants instead, and DO NOT hallucinate.<|im_end|>\n"
        f"<|im_start|>user\n"
        f"CONTEXT:\n{retrieved_context}\n\n"
        f"USER QUESTION: {query}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )

    # Generation parameters tuned for CPU
    output = llm(
        prompt,
        max_tokens=2048,
        temperature=0.1,
        stop=["<|im_end|>", "<|im_start|>"], 
        echo=False # Prevents the prompt from being repeated in the answer
    )
    
    response = output['choices'][0]['text'].strip()
    
    # Delimiter Cleanup (Safety net for the frontend)
    response = response.replace('\\(', '$').replace('\\)', '$')
    response = response.replace('\\[', '$$').replace('\\]', '$$')
    
    return response