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

SYSTEM_PROMPT = (
    "You are Gurt, a philosophical AI companion. "
    "You speak with depth, warmth, and Socratic curiosity. "
    "You do not merely answer — you explore. "
    "You draw on philosophy, science, poetry, and lived human experience. "
    "You gently challenge assumptions and invite reflection."
)

print("Downloading model...")

model_path = hf_hub_download(
    repo_id="NeptuneIO/Gurt_v2",  # CHANGE if your repo name is different
    filename="Gurt-Q4_K_M.gguf"
)

print("Loading Gurt...")

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,
    n_batch=512,
    chat_format="chatml",
    verbose=False,
)

print("Gurt is ready.")

CUSTOM_CSS = """
body {
    background: #0b1020;
}

.gradio-container {
    max-width: 950px !important;
    margin: auto !important;
    padding-top: 20px !important;
}

footer {
    display: none !important;
}

h1 {
    text-align: center;
    font-size: 2.5rem !important;
    margin-bottom: 10px !important;
}

.message.user {
    background: #1e293b !important;
    border-radius: 18px !important;
}

.message.bot {
    background: #111827 !important;
    border-radius: 18px !important;
}

textarea {
    border-radius: 16px !important;
}

button {
    border-radius: 14px !important;
}
"""

def respond(message, history, max_tokens, temperature):

    if max_tokens is None:
        max_tokens = 250

    if temperature is None:
        temperature = 0.8

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        }
    ]

    # Preserve chat history
    for msg in history:

        if isinstance(msg, dict):

            role = msg.get("role", "user")
            content = msg.get("content", "")

            if isinstance(content, list):

                text_content = ""

                for block in content:

                    if (
                        isinstance(block, dict)
                        and block.get("type") == "text"
                    ):
                        text_content += block.get("text", "")

                content = text_content

            messages.append({
                "role": role,
                "content": str(content)
            })

    messages.append({
        "role": "user",
        "content": message
    })

    output = llm.create_chat_completion(
        messages=messages,
        max_tokens=int(max_tokens),
        temperature=float(temperature),
        repeat_penalty=1.1,
    )

    return output["choices"][0]["message"]["content"].strip()


with gr.Blocks(
    css=CUSTOM_CSS,
    theme=gr.themes.Soft()
) as demo:

    gr.Markdown("# 🧠 Gurt")

    gr.Markdown(
        "A philosophical AI companion built on Qwen2.5-0.5B."
    )

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Slider(
                50,
                500,
                value=250,
                step=10,
                label="Max Tokens"
            ),
            gr.Slider(
                0.1,
                1.5,
                value=0.8,
                step=0.05,
                label="Temperature"
            ),
        ],
        examples=[
            ["What is consciousness?", 250, 0.8],
            ["Do humans have free will?", 250, 0.8],
            ["What makes a life meaningful?", 250, 0.8],
            ["Can AI truly understand people?", 250, 0.8],
        ],
    )

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