import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

from inverse_segmentation import GemmaInverseSegmenter

MODEL_ID = "dsba-lab/gemma2-9b-it-alienlm-full"
BASE_TOKENIZER_ID = "google/gemma-2-9b-it"
ALIEN_TOKENIZER_PATH = "assets/alien_tokenizer"

alien_tokenizer = AutoTokenizer.from_pretrained(
    ALIEN_TOKENIZER_PATH,
    use_fast=True,
)
base_tokenizer = AutoTokenizer.from_pretrained(
    BASE_TOKENIZER_ID,
    use_fast=True,
    token=os.environ.get("HF_TOKEN"),
)


def _vocab_id_range(tokenizer) -> tuple[int, int]:
    ids = tokenizer.get_vocab().values()
    return min(ids), max(ids)


def _assert_tokenizer_compatibility() -> None:
    """Apply the compatibility checks used by the official translator."""
    if alien_tokenizer.__class__.__name__ != base_tokenizer.__class__.__name__:
        raise RuntimeError("Alien and base tokenizer classes do not match.")
    if _vocab_id_range(alien_tokenizer) != _vocab_id_range(base_tokenizer):
        raise RuntimeError("Alien and base tokenizer ID ranges do not match.")
    if len(alien_tokenizer) != len(base_tokenizer):
        raise RuntimeError("Alien and base tokenizer vocabulary sizes do not match.")


_assert_tokenizer_compatibility()

# Client-side inverse of the fixed server tokenizer's text decoder. It receives
# only the decoded server string; generated token IDs are deliberately not
# exposed to this component.
inverse_segmenter = GemmaInverseSegmenter(base_tokenizer, alien_tokenizer)

# Load the model eagerly on CUDA
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    token=os.environ.get("HF_TOKEN"),
).to("cuda")
model.eval()


def plain2alien(text: str) -> str:
    token_ids = base_tokenizer.encode(text, add_special_tokens=False)
    return alien_tokenizer.decode(
        token_ids,
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )


def alien2plain(text: str) -> str:
    token_ids = alien_tokenizer.encode(text, add_special_tokens=False)
    return base_tokenizer.decode(
        token_ids,
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )


@spaces.GPU(duration=120)
def run_alienlm(
    user_prompt: str,
    max_new_tokens: int = 128,
    temperature: float = 0.7,
    top_p: float = 0.9,
    progress=gr.Progress(track_tqdm=True),
):
    """Run the full AlienLM privacy pipeline.

    1. Translate user's natural-language prompt into alien text.
    2. Feed the alienized prompt to the AlienLM-adapted model.
    3. Decode the model's alien-language response.
    4. Translate the alien response back into natural text.

    Args:
        user_prompt: Natural-language input from the user.
        max_new_tokens: Maximum number of new tokens to generate.
        temperature: Sampling temperature for generation.
        top_p: Nucleus-sampling probability threshold.
    """
    if not user_prompt.strip():
        return "", "", "", "Please enter a prompt."

    # Step 1: Alienize the user's prompt for display.
    alien_prompt = plain2alien(user_prompt)

    # Step 2: Match the SingleAlienTokenizer used for training: render the
    # plain chat, alienize the full rendered text, then encode it with the base
    # tokenizer. Special tokens are preserved by the tokenizer bijection.
    messages = [{"role": "user", "content": user_prompt}]
    plain_input_text = base_tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    alien_input_text = plain2alien(plain_input_text)
    inputs = base_tokenizer(
        alien_input_text,
        add_special_tokens=False,
        return_tensors="pt",
    ).to("cuda")

    # Step 3: Generate with the AlienLM model
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-5),
            top_p=top_p,
            pad_token_id=(
                alien_tokenizer.pad_token_id
                if alien_tokenizer.pad_token_id is not None
                else alien_tokenizer.eos_token_id
            ),
        )

    # Extract only the generated tokens (skip the prompt)
    generated_ids = output_ids[0][inputs["input_ids"].shape[1]:]
    # The training wrapper decodes model IDs with the base tokenizer first.
    alien_response = base_tokenizer.decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )

    # Step 4: Recover the hidden token segmentation from the server's text-only
    # response, then decode the selected ID path with the alien tokenizer. This
    # models a client that has no token IDs, logprobs, or server-side metadata.
    try:
        recovery = inverse_segmenter.recover(alien_response)
        recovered_response = recovery.text
    except ValueError:
        # Retain the official translator as a defensive fallback if an upstream
        # server applies an unexpected text normalization not present in Gemma.
        recovered_response = alien2plain(alien_response)

    return alien_prompt, alien_response, recovered_response, ""


# --- UI ---

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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🛸 AlienLM: Alienization of Language for API-Boundary Privacy\n"
            "Type a prompt in natural language. It gets translated into an *alien language* "
            "via a vocabulary-level token bijection, processed by an AlienLM-adapted Gemma-2-9B model, "
            "and the alien response is translated back to natural text — all without the model ever "
            "seeing your original words.\n\n"
            "📄 [Paper](https://arxiv.org/abs/2601.22710) · "
            "💻 [Code](https://github.com/KimJaehee0725/AlienLM) · "
            "🤗 [Model](https://huggingface.co/dsba-lab/gemma2-9b-it-alienlm-full)\n\n"
            "> **Demo scope:** This Space runs translation and inference in one backend process. "
            "In the intended deployment, the translator and bijection remain on the trusted client, "
            "so only alienized text crosses the model API boundary."
        )

        with gr.Row():
            prompt = gr.Textbox(
                label="Your prompt (natural language)",
                placeholder="e.g. Write a short poem about the ocean.",
                lines=3,
                scale=4,
            )
            run_btn = gr.Button("Run Pipeline", variant="primary", scale=1)

        with gr.Accordion("Advanced settings", open=False):
            max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max new tokens")
            temp = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
            top_p_val = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")

        with gr.Tabs():
            with gr.Tab("Recovered Response"):
                recovered_out = gr.Textbox(
                    label="Natural-language response (alien → plain)",
                    lines=8,
                    interactive=False,
                )
            with gr.Tab("Alienized Prompt"):
                alien_prompt_out = gr.Textbox(
                    label="Alienized prompt (what the model actually sees)",
                    lines=4,
                    interactive=False,
                )
            with gr.Tab("Alien Response"):
                alien_response_out = gr.Textbox(
                    label="Raw alien response from the model",
                    lines=8,
                    interactive=False,
                )

        error_box = gr.Textbox(label="Status", visible=False)

        gr.Examples(
            examples=[
                ["Write a short poem about the ocean."],
                ["Explain quantum computing to a child."],
                ["What is the meaning of life?"],
                ["Give me three tips for staying healthy."],
            ],
            inputs=[prompt],
            outputs=[alien_prompt_out, alien_response_out, recovered_out, error_box],
            fn=run_alienlm,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=run_alienlm,
        inputs=[prompt, max_tokens, temp, top_p_val],
        outputs=[alien_prompt_out, alien_response_out, recovered_out, error_box],
    )

    prompt.submit(
        fn=run_alienlm,
        inputs=[prompt, max_tokens, temp, top_p_val],
        outputs=[alien_prompt_out, alien_response_out, recovered_out, error_box],
    )


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