from __future__ import annotations

import spaces

import json
import os
import queue
import time
from pathlib import Path
from threading import Thread
from typing import Any, Generator

import gradio as gr
import torch
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)

from chat_utils import build_messages, estimate_duration

SETTINGS_PATH = Path(__file__).with_name("local_settings.json")
DEFAULTS = json.loads(SETTINGS_PATH.read_text(encoding="utf-8"))
MODEL_ID = os.getenv("MODEL_ID", DEFAULTS["model_id"])
MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", DEFAULTS["max_input_tokens"]))
HISTORY_MESSAGES = int(os.getenv("HISTORY_MESSAGES", DEFAULTS["history_messages"]))

print(f"Loading pre-quantized checkpoint: {MODEL_ID}")
processor = AutoProcessor.from_pretrained(MODEL_ID, padding_side="left")
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype="auto",
    low_cpu_mem_usage=True,
    attn_implementation="sdpa",
).eval()
model.to("cuda")
print("Quantized Gemma 4 model loaded.")


def _prepare_inputs(
    message: dict[str, Any],
    history: list[dict[str, Any]],
    system_prompt: str,
    thinking: bool,
) -> dict[str, torch.Tensor]:
    messages = build_messages(
        message,
        history,
        system_prompt,
        history_limit=HISTORY_MESSAGES,
    )
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
        enable_thinking=thinking,
    )
    input_tokens = int(inputs["input_ids"].shape[-1])
    if input_tokens > MAX_INPUT_TOKENS:
        raise gr.Error(
            f"Conversation uses {input_tokens:,} input tokens. "
            f"Clear or shorten it to at most {MAX_INPUT_TOKENS:,} tokens."
        )
    return {name: tensor.to(model.device) for name, tensor in inputs.items()}


@spaces.GPU(duration=estimate_duration, size="xlarge")
def chat(
    message: dict[str, Any],
    history: list[dict[str, Any]],
    system_prompt: str,
    thinking: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float,
) -> Generator[str, None, None]:
    """Generate a streamed response from typed text, text files, images, and chat history."""
    started = time.perf_counter()
    try:
        inputs = _prepare_inputs(message, history, system_prompt, thinking)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc

    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
        timeout=240,
    )
    generation_errors: queue.Queue[BaseException] = queue.Queue(maxsize=1)
    generated_token_counts: queue.Queue[int] = queue.Queue(maxsize=1)
    do_sample = float(temperature) > 0.05
    generation_kwargs: dict[str, Any] = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": int(max_new_tokens),
        "do_sample": do_sample,
        "repetition_penalty": float(repetition_penalty),
        "use_cache": True,
    }
    if do_sample:
        generation_kwargs.update(
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
        )

    def run_generation() -> None:
        try:
            with torch.inference_mode():
                outputs = model.generate(**generation_kwargs)
            generated_token_counts.put(
                max(0, int(outputs.shape[-1]) - int(inputs["input_ids"].shape[-1]))
            )
        except BaseException as exc:
            generation_errors.put(exc)
            streamer.on_finalized_text("", stream_end=True)

    worker = Thread(target=run_generation, daemon=True)
    worker.start()

    response = ""
    for piece in streamer:
        response += piece
        if response:
            yield response

    worker.join(timeout=1)
    if not generation_errors.empty():
        exc = generation_errors.get()
        raise gr.Error(f"Generation failed: {type(exc).__name__}: {exc}") from exc
    if not response.strip():
        yield "The model returned no visible text. Try another prompt or enable thinking."
    elif (
        not generated_token_counts.empty()
        and generated_token_counts.get() >= int(max_new_tokens)
    ):
        response += (
            "\n\n---\n"
            "*The response reached the output limit. Increase **Maximum new tokens** "
            "or ask the model to continue.*"
        )
        yield response

    elapsed = time.perf_counter() - started
    print(
        f"Generation complete: {elapsed:.1f}s, "
        f"input={inputs['input_ids'].shape[-1]}, output_limit={int(max_new_tokens)}"
    )


CSS = """
.gradio-container { max-width: 1180px !important; }
.hero { text-align: center; margin: 0 auto 0.75rem; }
.hero h1 { font-size: clamp(2rem, 5vw, 3.35rem); margin-bottom: 0.25rem; }
.hero p { color: var(--body-text-color-subdued); font-size: 1.05rem; }
.model-pill {
  display: inline-block; padding: .35rem .75rem; border-radius: 999px;
  background: color-mix(in srgb, var(--primary-500) 15%, transparent);
  border: 1px solid color-mix(in srgb, var(--primary-500) 35%, transparent);
}
"""

default_system = DEFAULTS["system_prompt"]
example_settings = [
    default_system,
    bool(DEFAULTS["thinking"]),
    int(DEFAULTS["max_new_tokens"]),
    float(DEFAULTS["temperature"]),
    float(DEFAULTS["top_p"]),
    int(DEFAULTS["top_k"]),
    float(DEFAULTS["repetition_penalty"]),
]

with gr.Blocks() as demo:
    gr.HTML(
        """
        <section class="hero">
          <div class="model-pill">FP8 · 31.3B parameters · text + vision</div>
          <h1>Gemma 4 Heretic</h1>
          <p>Refusal-reduced Gemma 4, served from a pre-quantized checkpoint.</p>
        </section>
        """
    )
    with gr.Accordion("Generation settings", open=False):
        system_prompt = gr.Textbox(
            value=default_system,
            label="System prompt",
            lines=2,
        )
        with gr.Row():
            thinking = gr.Checkbox(
                value=DEFAULTS["thinking"],
                label="Thinking mode",
                info="Can increase latency and output length.",
            )
            max_new_tokens = gr.Slider(
                128,
                1536,
                value=DEFAULTS["max_new_tokens"],
                step=128,
                label="Maximum new tokens",
                info="Longer answers use more of your ZeroGPU quota.",
            )
        with gr.Row():
            temperature = gr.Slider(
                0,
                1.5,
                value=DEFAULTS["temperature"],
                step=0.05,
                label="Temperature",
            )
            top_p = gr.Slider(
                0.1,
                1,
                value=DEFAULTS["top_p"],
                step=0.05,
                label="Top-p",
            )
            top_k = gr.Slider(
                1,
                128,
                value=DEFAULTS["top_k"],
                step=1,
                label="Top-k",
            )
            repetition_penalty = gr.Slider(
                1,
                1.3,
                value=DEFAULTS["repetition_penalty"],
                step=0.01,
                label="Repetition penalty",
            )

    chatbot = gr.Chatbot(
        height=590,
        layout="bubble",
        placeholder=(
            "<strong>Start a conversation</strong><br>"
            "Type a message or attach up to two images or UTF-8 text files."
        ),
        buttons=["copy", "copy_all"],
    )
    gr.ChatInterface(
        fn=chat,
        chatbot=chatbot,
        multimodal=True,
        textbox=gr.MultimodalTextbox(
            file_types=["image", "text"],
            file_count="multiple",
            placeholder="Type a message or attach image/text files...",
        ),
        additional_inputs=[
            system_prompt,
            thinking,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
            repetition_penalty,
        ],
        examples=[
            [
                {"text": "Explain why the sky is blue in three concise steps.", "files": []},
                *example_settings,
            ],
            [
                {"text": "Write a Python function that merges overlapping intervals.", "files": []},
                *example_settings,
            ],
            [
                {"text": "Give me a rigorous argument for and against open model weights.", "files": []},
                *example_settings,
            ],
        ],
        example_labels=["Explain", "Code", "Debate"],
        cache_examples=True,
        cache_mode="lazy",
        concurrency_limit=1,
        api_name="chat",
        api_description=chat.__doc__,
        flagging_mode="never",
        save_history=True,
    )
    gr.Markdown(
        """
        **Runtime note:** First use can take several minutes while the Space loads
        the 34 GB quantized checkpoint. ZeroGPU xlarge requests use twice the
        normal visitor quota. Conversations accept up to 16,384 input tokens and
        responses can use up to 1,536 new tokens. Text files must be UTF-8 and no
        larger than 64 KB. Output is unfiltered and can be wrong or unsafe.
        """
    )

demo.queue(default_concurrency_limit=1, max_size=12)

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="purple"),
        css=CSS,
        mcp_server=True,
    )