"""
Gradio ZeroGPU Space — Ollama-style LLM chat using llama-cpp-python.

GPU is acquired per request via @spaces.GPU, so the A100 is only held
while a token is being generated, not for the entire session lifetime.
"""

from __future__ import annotations

import os
import threading
from typing import Iterator

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

# ---------------------------------------------------------------------------
# Model configuration — change these to switch models
# ---------------------------------------------------------------------------
MODEL_REPO   = os.getenv("MODEL_REPO",   "LiquidAI/LFM2.5-230M-GGUF")
MODEL_FILE   = os.getenv("MODEL_FILE",   "LFM2.5-230M-F16.gguf")
CONTEXT_SIZE = int(os.getenv("CONTEXT_SIZE", "4096"))

# ---------------------------------------------------------------------------
# Load model once at startup (CPU map; GPU layers allocated at inference time)
# ---------------------------------------------------------------------------
print(f"Downloading {MODEL_FILE} from {MODEL_REPO} …")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
print("Model downloaded. Initialising llama-cpp …")

llm = Llama(
    model_path=model_path,
    n_ctx=CONTEXT_SIZE,
    n_gpu_layers=-1,   # offload ALL layers to GPU (set 0 for CPU-only)
    verbose=False,
)
print("Model ready.")

# ---------------------------------------------------------------------------
# Inference — wrapped with @spaces.GPU so A100 is acquired per call
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def _generate(
    messages: list[dict],
    temperature: float,
    max_new_tokens: int,
    top_p: float,
) -> Iterator[str]:
    """Yield partial assistant responses token by token."""
    stream = llm.create_chat_completion(
        messages=messages,
        temperature=temperature,
        max_tokens=max_new_tokens,
        top_p=top_p,
        stream=True,
    )
    for chunk in stream:
        delta = chunk["choices"][0]["delta"]
        token = delta.get("content", "")
        if token:
            yield token


def build_messages(
    history: list[dict],
    system_prompt: str,
) -> list[dict]:
    """Convert Gradio history format to llama-cpp messages list."""
    messages: list[dict] = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    return messages


def chat_fn(
    message: str,
    history: list[dict],
    system_prompt: str,
    temperature: float,
    max_new_tokens: int,
    top_p: float,
) -> Iterator[str]:
    """Gradio streaming chat handler."""
    history = history or []
    history.append({"role": "user", "content": message})
    messages = build_messages(history, system_prompt)

    partial = ""
    for token in _generate(messages, temperature, max_new_tokens, top_p):
        partial += token
        yield partial


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
DEFAULT_SYSTEM = (
    "You are a helpful, harmless, and honest AI assistant. "
    "Answer concisely and clearly."
)

with gr.Blocks(title="ZeroGPU LLM Chat", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
# 🦙 ZeroGPU LLM Chat
**Model:** `{MODEL_REPO} / {MODEL_FILE}`  
GPU is allocated on demand (ZeroGPU) — first response may take a few seconds while the Space warms up.
"""
    )

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.ChatInterface(
                fn=chat_fn,
                type="messages",
                additional_inputs_accordion=gr.Accordion(
                    label="⚙️ Generation settings", open=False
                ),
                additional_inputs=[
                    gr.Textbox(
                        value=DEFAULT_SYSTEM,
                        label="System prompt",
                        lines=3,
                        placeholder="Enter a system prompt …",
                    ),
                    gr.Slider(
                        minimum=0.0,
                        maximum=2.0,
                        value=0.7,
                        step=0.05,
                        label="Temperature",
                        info="Higher = more creative, lower = more deterministic",
                    ),
                    gr.Slider(
                        minimum=64,
                        maximum=2048,
                        value=512,
                        step=64,
                        label="Max new tokens",
                        info="Maximum number of tokens to generate per reply",
                    ),
                    gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.95,
                        step=0.05,
                        label="Top-p (nucleus sampling)",
                    ),
                ],
                examples=[
                    "Explain quantum entanglement in simple terms.",
                    "Write a Python function that checks if a string is a palindrome.",
                    "What are the pros and cons of renewable energy?",
                    "Translate 'Hello, how are you?' into French, German, and Japanese.",
                ],
                cache_examples=False,
            )

demo.queue(max_size=10)

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