import os
import gc
from typing import Iterator

# If running in environment without spaces, provide a no-op fallback for spaces.GPU
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(func=None, duration=None, size=None):
            if func is None:
                return lambda f: f
            return func

import gradio as gr
from huggingface_hub import hf_hub_download

MODEL_REPO = "huihui-ai/Huihui-Qwen3.8-27B-abliterated-GGUF"
# Default to UD-IQ4_XS which provides great balance of speed & quality
MODEL_FILE = "Huihui-Qwen3.8-27B-abliterated-UD-IQ4_XS.gguf"

print(f"Ensuring model {MODEL_FILE} is available...", flush=True)
MODEL_PATH = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
)
print(f"Model ready at: {MODEL_PATH}", flush=True)


def estimate_duration(
    history: list[dict[str, str]],
    system_prompt: str,
    temperature: float,
    top_p: float,
    top_k: int,
    max_tokens: int,
    repeat_penalty: float,
    *args,
    **kwargs,
) -> int:
    """Reserve a realistic ZeroGPU execution window based on requested max tokens."""
    tokens = int(max_tokens) if max_tokens else 512
    return min(180, max(30, int(tokens / 15) + 25))


def add_user_message(
    message: str,
    history: list[dict[str, str]],
) -> tuple[str, list[dict[str, str]]]:
    """Appends user message to chat history and clears input box."""
    if not message.strip():
        return "", history
    return "", history + [{"role": "user", "content": message.strip()}]


@spaces.GPU(duration=estimate_duration)
def bot_response(
    history: list[dict[str, str]],
    system_prompt: str,
    temperature: float,
    top_p: float,
    top_k: int,
    max_tokens: int,
    repeat_penalty: float,
) -> Iterator[list[dict[str, str]]]:
    """Streams the assistant's reply for the current conversation history.

    Args:
        history: Current conversation history including user's latest query.
        system_prompt: System prompt defining assistant persona.
        temperature: Sampling temperature (higher = more creative).
        top_p: Nucleus sampling probability cutoff.
        top_k: Top-K tokens to sample from.
        max_tokens: Maximum new tokens to generate.
        repeat_penalty: Penalty factor applied to repeated tokens.
    """
    if not history or history[-1].get("role") != "user":
        yield history
        return

    # Import llama_cpp inside the ZeroGPU worker so CUDA initializes in the GPU context
    from llama_cpp import Llama

    user_query = history[-1]["content"]
    prior_history = history[:-1]

    # Build chat messages sequence
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})

    for item in prior_history[-10:]:
        if isinstance(item, dict) and "role" in item and "content" in item:
            if item["content"]:
                messages.append({"role": item["role"], "content": item["content"]})

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

    # Prepare chat history with empty assistant bubble
    active_history = history + [{"role": "assistant", "content": ""}]
    yield active_history

    print("Initializing llama.cpp model on GPU...", flush=True)
    llm = Llama(
        model_path=MODEL_PATH,
        n_gpu_layers=-1,
        n_ctx=8192,
        n_batch=512,
        flash_attn=True,
        use_mmap=True,
        verbose=False,
    )

    try:
        response_stream = llm.create_chat_completion(
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            repeat_penalty=float(repeat_penalty),
            stream=True,
        )

        for chunk in response_stream:
            delta = chunk.get("choices", [{}])[0].get("delta", {})
            token = delta.get("content", "")
            if token:
                active_history[-1]["content"] += token
                yield active_history

    finally:
        del llm
        gc.collect()


CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# ⚡ Huihui Qwen3.8 27B Abliterated (GGUF)\n\n"
            "Fast conversational chat demo for "
            "[**huihui-ai/Huihui-Qwen3.8-27B-abliterated-GGUF**](https://huggingface.co/huihui-ai/Huihui-Qwen3.8-27B-abliterated-GGUF) "
            "powered by **llama.cpp** on Hugging Face **ZeroGPU**.\n\n"
            "> ⚠️ **Model Notice**: This is an uncensored / abliterated variant with reduced safety refusal filters. "
            "Outputs may contain sensitive or unfiltered responses. Use responsibly."
        )

        chatbot = gr.Chatbot(
            type="messages",
            height=540,
            layout="bubble",
            show_copy_button=True,
        )

        with gr.Row():
            message = gr.Textbox(
                placeholder="Ask anything or enter a prompt...",
                show_label=False,
                container=False,
                scale=5,
                autofocus=True,
            )
            send = gr.Button("Send", variant="primary", scale=1)

        with gr.Accordion("⚙️ Parameters & System Prompt", open=False):
            system_prompt = gr.Textbox(
                label="System Prompt",
                value="You are a helpful, precise, and honest AI assistant.",
                lines=2,
            )
            with gr.Row():
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
                top_k = gr.Slider(1, 100, value=40, step=1, label="Top-K")
            with gr.Row():
                max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max Tokens")
                repeat_penalty = gr.Slider(1.0, 1.5, value=1.1, step=0.05, label="Repetition Penalty")

        with gr.Row():
            clear = gr.ClearButton([message, chatbot], value="🗑️ Clear Chat")

        gr.Examples(
            examples=[
                ["Explain quantum computing in simple terms."],
                ["Write a fast Python script to parse and extract JSON data from nested API responses."],
                ["What are the key trade-offs between monolithic and microservice architectures?"],
                ["Compose a sci-fi short story about an AI discovering ancient human technology."],
            ],
            inputs=[message],
        )

        event_inputs = [
            chatbot,
            system_prompt,
            temperature,
            top_p,
            top_k,
            max_tokens,
            repeat_penalty,
        ]

        # Submit triggers user message display first, then streams assistant response
        message.submit(
            add_user_message,
            inputs=[message, chatbot],
            outputs=[message, chatbot],
            queue=False,
        ).then(
            bot_response,
            inputs=event_inputs,
            outputs=chatbot,
            api_name="chat",
        )

        send.click(
            add_user_message,
            inputs=[message, chatbot],
            outputs=[message, chatbot],
            queue=False,
        ).then(
            bot_response,
            inputs=event_inputs,
            outputs=chatbot,
            api_name="chat",
        )

        clear.click(lambda: [], outputs=chatbot, queue=False)

demo.queue(default_concurrency_limit=1)

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