import gradio as gr
import subprocess
from huggingface_hub import hf_hub_download

# 1. Install llama-cpp-python i runtime (inte via requirements.txt)
#    Viktigt: ta bort `llama-cpp-python` från requirements.txt,
#    annars försöker Spaces bygga från källkod och fastnar.
subprocess.run("pip install -q 'llama_cpp_python==0.3.15'", shell=True, check=False)
from llama_cpp import Llama


# 2. Ladda din GGUF-modell från Hugging Face
MODEL_REPO = "Jeppcode/ScalableLab2"
GGUF_FILENAME = "model-q4_k_m.gguf"  # eller "model-f16.gguf" om du vill ha fp16-varianten

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,      # kontextlängd
    n_threads=2,     # trådar (Spaces CPU är begränsad)
    n_batch=64,      # batchstorlek för generation
    use_mmap=True,
    use_mlock=False,
)


# 3. Några stil-lägen som "system prompts"
STYLE_SYSTEM_PROMPTS = {
    "Default": "You are a helpful, polite assistant.",
    "Short answer": (
        "You are a helpful assistant. Answer as concisely as possible, usually in 1–3 sentences."
    ),
    "Detailed explanation": (
        "You are a helpful teaching assistant. Give clear, structured and detailed explanations, "
        "often with bullet points or numbered steps when useful."
    ),
    "Step-by-step reasoning": (
        "You are a careful problem solver. Think step by step and explain your reasoning clearly "
        "before giving the final answer."
    ),
}


def _extract_text_from_content(content):
    """
    Gradio 6 ChatInterface använder 'messages'-format.
    content kan vara:
      - en sträng
      - en lista av blocks: [{"type": "text", "text": "..."} , ...]
    Vi konverterar det till en enkel sträng.
    """
    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, style):
    """
    Bygger en enkel textprompt för llama.cpp baserat på:
      - vald stil (system prompt)
      - konversationshistorik
      - senaste user-meddelandet
    Vi använder ett simpelt format:
        System: ...
        Conversation:
        User: ...
        Assistant: ...
        ...
        User: <current message>
        Assistant:
    """
    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])

    prompt_parts = []
    prompt_parts.append(f"System: {system_prompt}\n")
    prompt_parts.append("Conversation:\n")

    # history är en lista av dicts: {"role": "...", "content": ...}
    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")

    # Nuvarande användarmeddelande
    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, temperature, top_p, repetition_penalty, style):
    """
    Huvudfunktionen som Gradio ChatInterface anropar.
    - message: senaste user input
    - history: tidigare meddelanden (messages-format)
    - övriga parametrar: sliders / radio-knappar
    """
    prompt = build_prompt(message, history, style)

    # Hantera deterministiskt läge om temperature == 0
    temp = float(temperature)
    top_p_val = float(top_p)
    repeat_pen = float(repetition_penalty)

    if temp <= 0.0:
        temp = 0.0
        top_p_val = 1.0  # spelar mindre roll när temp=0

    output = llm(
        prompt,
        max_tokens=int(max_new_tokens),
        temperature=temp,
        top_p=top_p_val,
        repeat_penalty=repeat_pen,
        stop=["User:", "Assistant:", "System:", "Conversation:"],
    )

    reply = output["choices"][0]["text"].strip()
    return reply


# 4. DJ-reglagen (extra inputs till ChatInterface)
max_new_tokens_slider = gr.Slider(
    minimum=16,
    maximum=256,
    value=64,
    step=8,
    label="Max new tokens (response length)",
)

temperature_slider = gr.Slider(
    minimum=0.0,
    maximum=1.5,
    value=0.0,
    step=0.1,
    label="Temperature (0 = deterministic, higher = more random)",
)

top_p_slider = gr.Slider(
    minimum=0.1,
    maximum=1.0,
    value=0.9,
    step=0.05,
    label="Top-p (nucleus sampling)",
)

repetition_penalty_slider = gr.Slider(
    minimum=0.8,
    maximum=1.3,
    value=1.0,
    step=0.05,
    label="Repetition penalty",
)

style_radio = gr.Radio(
    choices=[
        "Default",
        "Short answer",
        "Detailed explanation",
        "Step-by-step reasoning",
    ],
    value="Detailed explanation",
    label="Answer style",
)

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.\n\n"
        "Use the controls in the accordion below like a DJ board to tweak "
        "response length, randomness and style."
    ),
    additional_inputs=[
        max_new_tokens_slider,
        temperature_slider,
        top_p_slider,
        repetition_penalty_slider,
        style_radio,
    ],
    additional_inputs_accordion="Generation controls",
)

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