import gradio as gr
import subprocess
from huggingface_hub import hf_hub_download

# 1. Install llama-cpp-python in runtime
subprocess.run("pip install -q 'llama_cpp_python==0.3.15'", shell=True, check=False)

from llama_cpp import Llama

# 2. Load GGUF model from Hugging Face
MODEL_REPO = "Jeppcode/ScalableLab2"
GGUF_FILENAME = "model-q4_k_m.gguf"

print(f"Downloading GGUF model {MODEL_REPO}/{GGUF_FILENAME} ...")
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=GGUF_FILENAME,
)

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,
)

# Hardcoded system prompt (since style selector is removed)
SYSTEM_PROMPT = "You are a helpful teaching assistant. Give clear, structured and detailed explanations."

def _extract_text_from_content(content):
    """
    Extracts text from Gradio 6 message format.
    """
    if isinstance(content, list):
        texts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                texts.append(block.get("text", ""))
            else:
                texts.append(str(block))
        return "\n".join(t for t in texts if t)
    else:
        return str(content)

def build_prompt(message, history):
    """
    Builds the prompt using the hardcoded system prompt.
    """
    prompt_parts = []
    prompt_parts.append(f"System: {SYSTEM_PROMPT}\n")
    prompt_parts.append("Conversation:\n")
    
    for msg in history or []:
        role = msg.get("role")
        content = _extract_text_from_content(msg.get("content", ""))
        if not content:
            continue
        if role == "user":
            prompt_parts.append(f"User: {content}\n")
        elif role == "assistant":
            prompt_parts.append(f"Assistant: {content}\n")
        elif role == "system":
            prompt_parts.append(f"System (previous): {content}\n")
            
    # Current user message
    prompt_parts.append(f"User: {message}\n")
    prompt_parts.append("Assistant:")
    full_prompt = "".join(prompt_parts)
    return full_prompt

def chat_fn(message, history, max_new_tokens):
    """
    Main chat function. 
    Only accepts max_new_tokens as an additional input now.
    """
    prompt = build_prompt(message, history)
    
    # Internal defaults for the removed sliders
    temperature = 0.7
    top_p = 0.9
    repetition_penalty = 1.0

    output = llm(
        prompt,
        max_tokens=int(max_new_tokens),
        temperature=temperature,
        top_p=top_p,
        repeat_penalty=repetition_penalty,
        stop=["User:", "Assistant:", "System:", "Conversation:"],
    )
    reply = output["choices"][0]["text"].strip()
    return reply

# 4. Only keep the Max Token slider
max_new_tokens_slider = gr.Slider(
    minimum=16,
    maximum=256,
    value=64,
    step=8,
    label="Max new tokens (response length)",
)

demo = gr.ChatInterface(
    fn=chat_fn,
    title="Lab 2 – Fine-tuned GGUF model",
    description=(
        "Chat with our fine-tuned Llama-based model, converted to GGUF and "
        "loaded via llama.cpp from Jeppcode/ScalableLab2."
    ),
    additional_inputs=[
        max_new_tokens_slider,
    ],
    additional_inputs_accordion="Generation controls",
)

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