import gradio as gr
import subprocess
from huggingface_hub import hf_hub_download

# --- 1. Setup & Install ---
subprocess.run("pip install -q 'llama_cpp_python==0.3.15'", shell=True, check=False)
from llama_cpp import Llama

# --- 2. Load Model (GGUF) ---
MODEL_REPO = "Jeppcode/ScalableLab2"
GGUF_FILENAME = "model-q4_k_m.gguf"

print(f"Downloading GGUF model {MODEL_REPO}/{GGUF_FILENAME} ...")
try:
    model_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=GGUF_FILENAME,
    )
except Exception as e:
    print(f"Error downloading model: {e}")
    model_path = ""

llm = None
if model_path:
    print("Initializing llama.cpp LLM ...")
    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        n_threads=2,
        n_batch=64,
        use_mmap=True,
        use_mlock=False,
    )

# --- 3. Style / System Prompts ---
# These are the "Buttons" logic to change how the AI behaves
STYLE_SYSTEM_PROMPTS = {
    "Default": (
        "You are a helpful, polite AI assistant. "
    ),
    "Short answer": (
        "Answer as briefly as possible, usually in 1-3 sentences. "
        "Give only the core information needed to answer the question. "
        "Do not add extra explanations, lists, or examples unless the user asks for more detail."
    ),
    "Detailed explanation": (
        "Give a clear, structured, and detailed explanation. "
        "Break your answer into short paragraphs or bullet points when helpful. "
        "Explain what, how, and why, but avoid unnecessary repetition or filler."
    ),
    "Step-by-step reasoning": (
        "Solve the problem step by step. "
        "First restate the task in your own words, then explain your reasoning in numbered steps, "
        "and finally give a short final answer at the end. "
        "Keep the reasoning easy to follow and avoid unrelated digressions."
    ),
}

def _extract_text(content):
    if isinstance(content, list):
        return "\n".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
    return str(content)

def chat_fn(message, history, max_new_tokens, style):
    if not llm: return "Error: Model not loaded."
    
    # Select the specific system prompt based on the button chosen
    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
    
    prompt = f"System: {system_prompt}\nConversation:\n"
    for msg in history or []:
        role = msg.get("role")
        txt = _extract_text(msg.get("content", ""))
        if txt:
            if role == "user": prompt += f"User: {txt}\n"
            elif role == "assistant": prompt += f"Assistant: {txt}\n"
            
    prompt += f"User: {message}\nAssistant:"

    # Default internal values for randomness
    output = llm(
        prompt,
        max_tokens=int(max_new_tokens),
        temperature=0.7,
        top_p=0.9,
        stop=["User:", "Assistant:", "System:"],
    )
    return output["choices"][0]["text"].strip()

# --- 4. UI Controls ---

# Slider for length
max_new_tokens_slider = gr.Slider(
    minimum=16, 
    maximum=256, 
    value=64, 
    step=8, 
    label="Max Response Length"
)

# The "Buttons" at the bottom for Style
style_radio = gr.Radio(
    choices=["Default", "Short answer", "Detailed explanation", "Step-by-step reasoning"],
    value="Detailed explanation",
    label="Answer Style"
)

# --- 5. Launch App (Clean / No Theme) ---
demo = gr.ChatInterface(
    fn=chat_fn,
    title="Lab 2 – Fine-tuned GGUF model",
    description="Chat with the fine-tuned Llama model. Use the controls below to change the response style.",
    additional_inputs=[max_new_tokens_slider, style_radio],
    additional_inputs_accordion="Controls",
)

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