import re
import spaces
import torch
import torch.nn.functional as F
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModel

MODEL_ID = "GSAI-ML/iLLaDA-8B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(
    MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16
).to("cuda").eval()

MASK_TOKEN = "[MASK]"
MASK_ID = 5  # iLLaDA <[MASK]> token id


def extract_text(content):
    """Extract plain text from Gradio message content (may be str or list of dicts)."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            item.get("text", "") if isinstance(item, dict) else str(item)
            for item in content
        )
    return str(content)


def split_thinking(text):
    """Split an iLLaDA response into (thinking, answer).
    Returns (None, text) when there is no <think> block."""
    m = re.search(r"<think>(.*?)</think>(.*)", text, re.DOTALL)
    if m:
        return m.group(1).strip(), m.group(2).strip()
    if "<think>" in text:  # not yet closed
        return text.split("<think>", 1)[1].strip(), ""
    return None, text.strip()


def parse_constraints(constraints_text):
    """Parse 'position:word, position:word, ...' into a dict mapping
    gen-relative token positions to token IDs."""
    constraints = {}
    if not constraints_text:
        return constraints
    for part in constraints_text.split(","):
        if ":" not in part:
            continue
        pos_str, word = part.split(":", 1)
        try:
            pos = int(pos_str.strip())
        except ValueError:
            continue
        word = word.strip()
        if not word or pos < 0:
            continue
        token_ids = tokenizer.encode(" " + word, add_special_tokens=False)
        for i, tid in enumerate(token_ids):
            constraints[pos + i] = tid
    return constraints


def confidence_label(prob):
    if prob < 0.3:
        return "low"
    elif prob < 0.7:
        return "mid"
    return "high"


def build_vis_state(x, prompt_length, gen_length, confidences=None):
    """Build (highlighted_text_state, plain_text) from the current token sequence."""
    highlighted = []
    tokens = []
    for i in range(gen_length):
        pos = prompt_length + i
        if pos >= x.shape[1] or x[0, pos].item() == MASK_ID:
            highlighted.append((MASK_TOKEN, None))
            tokens.append(MASK_TOKEN)
        else:
            token = tokenizer.decode([x[0, pos].item()], skip_special_tokens=True)
            token = token or " "
            label = confidence_label(confidences[i]) if confidences and i in confidences else None
            highlighted.append((token, label))
            tokens.append(token)
    return highlighted, "".join(tokens)


def add_gumbel_noise(logits, temperature):
    """Gumbel-max sampling in float64 (as in the official LLaDA implementation)."""
    if temperature == 0:
        return logits
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    gumbel_noise = (-torch.log(noise)) ** temperature
    return logits.exp() / gumbel_noise


def get_num_transfer_tokens(mask_index, steps):
    """How many tokens to un-mask at each step so masks deplete evenly."""
    mask_num = mask_index.sum(dim=1, keepdim=True)
    base = mask_num // steps
    remainder = mask_num % steps
    num_transfer_tokens = (
        torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64)
        + base
    )
    for i in range(mask_num.size(0)):
        num_transfer_tokens[i, : remainder[i]] += 1
    return num_transfer_tokens


@spaces.GPU
@torch.no_grad()
def generate_streaming(messages, gen_length, steps, temperature, block_length,
                       cfg_scale, remasking, constraints=None):
    """Streaming semi-autoregressive diffusion generation.
    Yields (highlighted, plain, response_text_or_None) at each denoising step.
    """
    if constraints is None:
        constraints = {}

    prompt_text = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    input_ids = tokenizer(prompt_text)["input_ids"]
    prompt = torch.tensor(input_ids, device="cuda").unsqueeze(0)
    prompt_length = prompt.shape[1]

    # gen_length must be divisible by block_length
    gen_length = max(block_length, (gen_length // block_length) * block_length)
    num_blocks = gen_length // block_length
    # total steps split across blocks
    steps = max(num_blocks, (steps // num_blocks) * num_blocks)
    steps_per_block = steps // num_blocks

    x = torch.full((1, prompt_length + gen_length), MASK_ID, dtype=torch.long, device="cuda")
    x[:, :prompt_length] = prompt.clone()

    # Pin constrained tokens into the initial sequence (treated as fixed context)
    for gen_pos, tid in constraints.items():
        abs_pos = prompt_length + gen_pos
        if abs_pos < x.shape[1]:
            x[0, abs_pos] = tid

    prompt_index = x != MASK_ID

    token_confidences = {gen_pos: 1.0 for gen_pos in constraints if 0 <= gen_pos < gen_length}
    highlighted, plain = build_vis_state(x, prompt_length, gen_length, token_confidences)
    yield highlighted, plain, None

    for num_block in range(num_blocks):
        block_start = prompt_length + num_block * block_length
        block_end = prompt_length + (num_block + 1) * block_length
        block_mask_index = x[:, block_start:block_end] == MASK_ID
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps_per_block)

        for i in range(steps_per_block):
            mask_index = x == MASK_ID

            if cfg_scale > 0.0:
                un_x = x.clone()
                un_x[prompt_index] = MASK_ID
                x_ = torch.cat([x, un_x], dim=0)
                logits = model(x_).logits
                logits, un_logits = torch.chunk(logits, 2, dim=0)
                logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
            else:
                logits = model(x).logits

            logits_with_noise = add_gumbel_noise(logits, temperature)
            x0 = torch.argmax(logits_with_noise, dim=-1)

            if remasking == "low_confidence":
                p = F.softmax(logits.to(torch.float64), dim=-1)
                x0_p = torch.gather(p, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1)
            else:  # random
                x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)

            # never un-mask beyond the current block
            x0_p[:, block_end:] = -np.inf

            x0 = torch.where(mask_index, x0, x)
            confidence = torch.where(mask_index, x0_p, -np.inf)

            transfer_index = torch.zeros_like(x0, dtype=torch.bool)
            for j in range(confidence.shape[0]):
                k = int(num_transfer_tokens[j, i])
                if k > 0:
                    _, select_index = torch.topk(confidence[j], k=k)
                    transfer_index[j, select_index] = True
            x[transfer_index] = x0[transfer_index]

            # record confidence of newly committed tokens
            for pos in transfer_index[0].nonzero(as_tuple=True)[0].tolist():
                gen_pos = pos - prompt_length
                if 0 <= gen_pos < gen_length:
                    if remasking == "low_confidence":
                        token_confidences[gen_pos] = float(x0_p[0, pos].item())
                    else:
                        token_confidences[gen_pos] = 1.0

            highlighted, plain = build_vis_state(x, prompt_length, gen_length, token_confidences)
            yield highlighted, plain, None

    generated = x[:, prompt_length:]
    response_text = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
    highlighted, plain = build_vis_state(x, prompt_length, gen_length, token_confidences)
    yield highlighted, plain, response_text


css = """
.category-legend{display:none}
button{height: 60px}
.legend{margin-bottom: 5px}
.legend-item{height: 25px}
"""


def create_chatbot_demo():
    with gr.Blocks() as demo:
        gr.Markdown("# iLLaDA-8B-Instruct - Masked Diffusion Language Model Demo")
        gr.Markdown(
            "[model iLLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/iLLaDA-8B-Instruct), "
            "[paper](https://arxiv.org/abs/2606.25331), "
            "[code](https://github.com/ML-GSAI/LLaDA)"
        )

        with gr.Row():
            with gr.Column(scale=3):
                chatbot_ui = gr.Chatbot(label="Conversation", height=500)

                with gr.Group():
                    with gr.Row():
                        user_input = gr.Textbox(
                            label="Your Message",
                            placeholder="Type your message here...",
                            show_label=False,
                        )
                        send_btn = gr.Button("Send")

                constraints_input = gr.Textbox(
                    label="Word Constraints",
                    info="Pin specific words at specific generated-token positions: "
                         "'position:word' format. Example: '0:Once, 5:upon, 10:time'",
                    placeholder="0:Once, 5:upon, 10:time",
                    value="",
                )

            with gr.Column(scale=2):
                output_vis = gr.HighlightedText(
                    label="Diffusion process (token confidence)",
                    combine_adjacent=False,
                    show_legend=True,
                    color_map={
                        "low": "#FF6666",
                        "mid": "#FFAA33",
                        "high": "#66CC66",
                    },
                )

        with gr.Accordion("Generation Settings", open=False):
            with gr.Row():
                gen_length = gr.Slider(
                    minimum=16, maximum=1024, value=256, step=8,
                    label="Generation Length",
                )
                steps = gr.Slider(
                    minimum=8, maximum=1024, value=256, step=8,
                    label="Denoising Steps",
                )
            with gr.Row():
                temperature = gr.Slider(
                    minimum=0.0, maximum=1.0, value=0.0, step=0.1,
                    label="Temperature",
                )
                block_length = gr.Slider(
                    minimum=8, maximum=128, value=32, step=8,
                    label="Block Length",
                )
            with gr.Row():
                cfg_scale = gr.Slider(
                    minimum=0.0, maximum=4.0, value=0.0, step=0.1,
                    label="CFG Scale (classifier-free guidance)",
                )
                remasking = gr.Radio(
                    choices=["low_confidence", "random"],
                    value="low_confidence",
                    label="Remasking Strategy",
                )

        clear_btn = gr.Button("Clear Conversation")

        def user_message_submitted(message, history):
            if not message.strip():
                return history, "", []
            history = history + [{"role": "user", "content": message}]
            return history, "", []

        def bot_response(history, gen_length, steps, temperature, block_length,
                         cfg_scale, remasking, constraints_text):
            if not history:
                yield history, []
                return

            try:
                messages = [
                    {"role": msg["role"], "content": extract_text(msg["content"])}
                    for msg in history
                    if msg["role"] in ("user", "assistant")
                    and msg.get("content")
                    and not (msg.get("metadata") or {}).get("title")
                ]

                constraints = parse_constraints(constraints_text)

                # Live diffusion shown inside a collapsible "thinking" panel.
                base = history
                history = base + [{
                    "role": "assistant",
                    "content": "",
                    "metadata": {"title": "💭 Diffusion process", "status": "pending"},
                }]

                response_text = None
                for vis_state, plain, text in generate_streaming(
                    messages, gen_length, steps, temperature, block_length,
                    cfg_scale, remasking, constraints,
                ):
                    if text is not None:
                        response_text = text
                    history[-1]["content"] = plain
                    yield history, vis_state

                # Final: split <think> reasoning from the answer.
                final = response_text if response_text is not None else history[-1]["content"]
                thinking, answer = split_thinking(final)
                new_msgs = []
                if thinking:
                    new_msgs.append({
                        "role": "assistant",
                        "content": thinking,
                        "metadata": {"title": "💭 Thought", "status": "done"},
                    })
                    new_msgs.append({
                        "role": "assistant",
                        "content": answer or "_(Reasoning was cut off — raise Generation "
                                            "Length / Denoising Steps for a full answer.)_",
                    })
                else:
                    new_msgs.append({"role": "assistant", "content": answer or final})
                history = base + new_msgs
                yield history, vis_state

            except Exception as e:
                error_msg = f"Error: {str(e)}"
                print(error_msg)
                yield history, [(error_msg, "low")]

        def clear_conversation():
            return [], "", []

        clear_btn.click(
            fn=clear_conversation,
            inputs=[],
            outputs=[chatbot_ui, user_input, output_vis],
        )

        msg_submit = user_input.submit(
            fn=user_message_submitted,
            inputs=[user_input, chatbot_ui],
            outputs=[chatbot_ui, user_input, output_vis],
        )
        send_click = send_btn.click(
            fn=user_message_submitted,
            inputs=[user_input, chatbot_ui],
            outputs=[chatbot_ui, user_input, output_vis],
        )

        bot_inputs = [
            chatbot_ui, gen_length, steps, temperature,
            block_length, cfg_scale, remasking, constraints_input,
        ]
        bot_outputs = [chatbot_ui, output_vis]

        msg_submit.then(fn=bot_response, inputs=bot_inputs, outputs=bot_outputs)
        send_click.then(fn=bot_response, inputs=bot_inputs, outputs=bot_outputs)

    return demo


if __name__ == "__main__":
    demo = create_chatbot_demo()
    demo.queue().launch(css=css)