import os
from pathlib import Path

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

MODEL_ID = os.getenv("MODEL", "LSX-UniWue/LLaMmlein_1B_chat_selected")
MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "LSX-UniWue/LLaMmlein_1B_alternative_formats")
MODEL_REVISION = os.getenv("MODEL_REVISION", "7d97b69ae6910b5f317be2dbd5b4820d848c66b4")
MODEL_FILENAME = os.getenv("MODEL_FILENAME", "LLaMmlein_1B_chat_selected.gguf")
QUANT = os.getenv("QUANT", "Q8_0/BF16 GGUF")
N_CTX = int(os.getenv("N_CTX", "2048"))
N_THREADS = int(os.getenv("N_THREADS", str(min(4, os.cpu_count() or 2))))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "64"))

model_name = MODEL_ID.split("/")[-1]
title = f"🇩🇪 {model_name}"
description = (
    f"Chat with <a href=\"https://huggingface.co/{MODEL_ID}\">{model_name}</a> "
    f"in GGUF format ({QUANT})."
)

print("resolving model file", flush=True)
model_path = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename=MODEL_FILENAME,
    revision=MODEL_REVISION,
)
print(f"loading model from {Path(model_path).name}", flush=True)

llm = Llama(
    model_path=model_path,
    n_ctx=N_CTX,
    n_threads=N_THREADS,
    n_batch=64,
    verbose=False,
)

def iter_history_messages(history):
    for entry in history or []:
        if isinstance(entry, dict):
            role = entry.get("role")
            content = entry.get("content")
            if role in {"user", "assistant"} and content:
                yield {"role": role, "content": content}
        else:
            human, assistant = entry
            if human:
                yield {"role": "user", "content": human}
            if assistant:
                yield {"role": "assistant", "content": assistant}

def chat_stream_completion(message, history):
    messages_prompts = list(iter_history_messages(history))
    messages_prompts.append({"role": "user", "content": message})
    prompt_parts = [
        "Du bist ein hilfreicher deutschsprachiger Assistent.",
        "",
    ]
    for item in messages_prompts:
        label = "Benutzer" if item["role"] == "user" else "Assistent"
        prompt_parts.append(f"{label}: {item['content']}")
    prompt_parts.append("Assistent:")
    prompt = "\n".join(prompt_parts)

    response = llm.create_completion(
        prompt=prompt,
        max_tokens=MAX_TOKENS,
        repeat_penalty=1.1,
        stream=True,
        stop=["\nBenutzer:", "\nAssistent:"],
    )
    message_repl = ""
    for chunk in response:
        text = chunk["choices"][0].get("text", "")
        if text:
            message_repl = message_repl + text
        yield message_repl

print("starting gradio", flush=True)
gr.ChatInterface(
    fn=chat_stream_completion,
    type="messages",
    title=title,
    description=description,
    examples=[
        ["Was weißt du über Würzburg?"],
        ["Erkläre Quantencomputing in einfachen Worten."],
    ],
    cache_examples=False,
).queue().launch()