from __future__ import annotations

import os
import random
import re
import sys
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Generator

import gradio as gr
import torch
from huggingface_hub import snapshot_download


MODEL_ID = "MarkChenX/lfm2-quantum-128m-sft-v2-reasoning"
HF_TOKEN = os.getenv("HF_TOKEN")

MAX_CONTEXT_TOKENS = 1024

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)


# ---------------------------------------------------------------------
# Download repository
# ---------------------------------------------------------------------

print(f"Downloading model repository: {MODEL_ID}")

model_dir = Path(
    snapshot_download(
        repo_id=MODEL_ID,
        token=HF_TOKEN,
        ignore_patterns=[
            "optim_*.pt",
            "optimizer*.pt",
            "*.bin",
        ],
    )
)

print(f"Model repository: {model_dir}")

# Make the bundled nanochat package importable.
sys.path.insert(0, str(model_dir))


# ---------------------------------------------------------------------
# Import custom architecture
# ---------------------------------------------------------------------

try:
    from nanochat.checkpoint_manager import build_model
    from nanochat.engine import Engine
except ImportError as error:
    raise RuntimeError(
        "Could not import the bundled NanoChat code. The model repository "
        "must contain nanochat/checkpoint_manager.py and nanochat/engine.py."
    ) from error


# ---------------------------------------------------------------------
# Locate latest model checkpoint
# ---------------------------------------------------------------------

def find_latest_checkpoint(
    directory: Path,
) -> tuple[Path, int]:
    candidates: list[tuple[int, Path]] = []

    for path in directory.rglob("model_*.pt"):
        match = re.fullmatch(
            r"model_(\d+)\.pt",
            path.name,
        )

        if match:
            candidates.append(
                (int(match.group(1)), path)
            )

    if not candidates:
        raise FileNotFoundError(
            f"No model_XXXXXX.pt checkpoint found under {directory}"
        )

    candidates.sort(
        key=lambda item: item[0],
        reverse=True,
    )

    step, checkpoint_path = candidates[0]
    return checkpoint_path, step


checkpoint_path, checkpoint_step = find_latest_checkpoint(model_dir)
checkpoint_dir = checkpoint_path.parent

metadata_path = checkpoint_dir / f"meta_{checkpoint_step:06d}.json"

if not metadata_path.exists():
    raise FileNotFoundError(
        "Checkpoint metadata is missing. Expected: "
        f"{metadata_path}"
    )


print(f"Checkpoint directory: {checkpoint_dir}")
print(f"Checkpoint file: {checkpoint_path}")
print(f"Checkpoint step: {checkpoint_step}")
print(f"Metadata file: {metadata_path}")
print(f"Loading on device: {device}")


# ---------------------------------------------------------------------
# Load custom NanoChat model
# ---------------------------------------------------------------------

model, tokenizer, metadata = build_model(
    str(checkpoint_dir),
    checkpoint_step,
    device,
    "eval",
)

model.eval()

engine = Engine(model, tokenizer)

assistant_end_token = tokenizer.encode_special(
    "<|assistant_end|>"
)

bos_token = tokenizer.get_bos_token_id()


if device.type == "cuda":
    autocast_context = lambda: torch.amp.autocast(
        device_type="cuda",
        dtype=torch.bfloat16,
    )
else:
    autocast_context = nullcontext


print("Model and NanoChat engine loaded successfully.")


# ---------------------------------------------------------------------
# Gradio history conversion
# ---------------------------------------------------------------------

def normalize_history(
    history: list[Any] | None,
) -> list[dict[str, str]]:
    messages: list[dict[str, str]] = []

    for item in history or []:
        # Current Gradio message format.
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")

            if (
                role in {"user", "assistant"}
                and isinstance(content, str)
                and content.strip()
            ):
                messages.append(
                    {
                        "role": role,
                        "content": content.strip(),
                    }
                )

        # Compatibility with older tuple-style Gradio history.
        elif isinstance(item, (tuple, list)) and len(item) == 2:
            user_content, assistant_content = item

            if isinstance(user_content, str) and user_content.strip():
                messages.append(
                    {
                        "role": "user",
                        "content": user_content.strip(),
                    }
                )

            if (
                isinstance(assistant_content, str)
                and assistant_content.strip()
            ):
                messages.append(
                    {
                        "role": "assistant",
                        "content": assistant_content.strip(),
                    }
                )

    return messages


# ---------------------------------------------------------------------
# Build NanoChat-native conversation tokens
# ---------------------------------------------------------------------

def build_conversation_tokens(
    message: str,
    history: list[Any] | None,
    system_message: str,
) -> list[int]:
    messages = normalize_history(history)

    current_message = message.strip()

    # NanoChat merges system instructions into the first user message.
    if system_message.strip():
        current_message = (
            system_message.strip()
            + "\n\n"
            + current_message
        )

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

    tokens: list[int] = [bos_token]

    user_start = tokenizer.encode_special("<|user_start|>")
    user_end = tokenizer.encode_special("<|user_end|>")

    assistant_start = tokenizer.encode_special(
        "<|assistant_start|>"
    )
    assistant_end = tokenizer.encode_special(
        "<|assistant_end|>"
    )

    for chat_message in messages:
        content_tokens = tokenizer.encode(
            chat_message["content"]
        )

        if chat_message["role"] == "user":
            tokens.append(user_start)
            tokens.extend(content_tokens)
            tokens.append(user_end)

        elif chat_message["role"] == "assistant":
            tokens.append(assistant_start)
            tokens.extend(content_tokens)
            tokens.append(assistant_end)

    # Prime the model to produce the next assistant message.
    tokens.append(assistant_start)

    # Leave room for generation within the model's 1024-token context.
    return tokens


# ---------------------------------------------------------------------
# Gradio generation function
# ---------------------------------------------------------------------

def respond(
    message: str,
    history: list[Any],
    system_message: str,
    max_new_tokens: int,
    temperature: float,
    top_k: int,
) -> Generator[str, None, None]:
    if not message or not message.strip():
        yield "Please enter a message."
        return

    max_new_tokens = int(max_new_tokens)
    temperature = float(temperature)
    top_k = int(top_k)

    try:
        prompt_tokens = build_conversation_tokens(
            message=message,
            history=history,
            system_message=system_message,
        )

        # Preserve enough context space for newly generated tokens.
        maximum_prompt_length = (
            MAX_CONTEXT_TOKENS - max_new_tokens
        )

        if maximum_prompt_length <= 0:
            yield (
                "Max new tokens must be smaller than the "
                f"{MAX_CONTEXT_TOKENS}-token context window."
            )
            return

        if len(prompt_tokens) > maximum_prompt_length:
            # Keep BOS and the newest conversation context.
            prompt_tokens = [
                bos_token,
                *prompt_tokens[-(maximum_prompt_length - 1):],
            ]

        generated_tokens: list[int] = []
        last_text = ""

        with autocast_context():
            stream = engine.generate(
                prompt_tokens,
                num_samples=1,
                max_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                seed=random.randint(0, 2**31 - 1),
            )

            for token_column, _token_masks in stream:
                token = token_column[0]

                if token in {
                    assistant_end_token,
                    bos_token,
                }:
                    break

                generated_tokens.append(token)

                current_text = tokenizer.decode(
                    generated_tokens
                )

                # Avoid yielding incomplete UTF-8 replacement characters.
                if current_text.endswith("�"):
                    continue

                if current_text != last_text:
                    last_text = current_text
                    yield current_text

        if not generated_tokens:
            yield "The model generated an empty response."

    except Exception as error:
        yield (
            "Generation failed: "
            f"{type(error).__name__}: {error}"
        )


# ---------------------------------------------------------------------
# Gradio interface
# ---------------------------------------------------------------------
with gr.Blocks(
    title="LFM2 Quantum 128M",
) as demo:
    gr.Markdown(
        """
        # LFM2 Quantum 128M

        Custom LFM2-style hybrid quantum language model running with
        the bundled NanoChat checkpoint code.
        """
    )

    gr.ChatInterface(
        fn=respond,
        chatbot=gr.Chatbot(
            height=520,
            placeholder="Ask the model a question.",
        ),
        additional_inputs=[
            gr.Textbox(
                value=(
                    "You are a helpful, concise, and friendly assistant."
                ),
                label="System instruction",
                lines=3,
            ),
            gr.Slider(
                minimum=16,
                maximum=512,
                value=256,
                step=16,
                label="Max new tokens",
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                value=0.7,
                step=0.1,
                label="Temperature",
            ),
            gr.Slider(
                minimum=1,
                maximum=200,
                value=50,
                step=1,
                label="Top-k",
            ),
        ],
        examples=[
            [
                "Hello! Introduce yourself.",
                "You are a helpful, concise, and friendly assistant.",
                256,
                0.7,
                50,
            ],
            [
                "Explain quantum machine learning simply.",
                "You are a helpful AI research assistant.",
                256,
                0.7,
                50,
            ],
        ],
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,
    )