import queue
import threading
import time

import gradio as gr
import spaces
import torch
from transformers import PreTrainedTokenizerFast

import mfn_hf  # shared look + usage analytics; import before @spaces.GPU

from flint import generate, load_flint

MODEL_ID = "RecursiveSI/flint"
TITLE = "Flint"
FP8_EXPERTS = False  # bf16 fits on a full RTX Pro 6000; FP8 experts are the fallback

# Prompt + reply share this window. Measured on an RTX Pro 6000 with the bf16 weights:
# 32K tokens peak at ~79 GiB of the 95 GiB card.
MAX_CONTEXT = 32768
MAX_NEW_TOKENS = 8192
# Measured: prefill ~12K tok/s at 8K-32K, decode 100-113 tok/s (CUDA-graph decode). Conservative.
PREFILL_TOK_S = 6000.0
DECODE_TOK_S = 60.0
MAX_GPU_SECONDS = 300

tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_ID)
model, _cfg = load_flint(MODEL_ID, device="cuda", dtype=torch.bfloat16, fp8_experts=FP8_EXPERTS, delete_shards=True)

EOS_IDS = {150001, 150003, 150005, 150018}
DEFAULT_SYSTEM = ""


def gpu_seconds(prompt_tokens, max_new_tokens):
    est = 10 + prompt_tokens / PREFILL_TOK_S + int(max_new_tokens) / DECODE_TOK_S
    return int(min(MAX_GPU_SECONDS, est))


def gpu_duration(message, history, system_prompt, max_new_tokens, *args):
    try:
        system, msgs = build_messages(message, history, system_prompt)
        n = int(encode(system, msgs, max_new_tokens).shape[0])
    except Exception:
        n = MAX_CONTEXT - int(max_new_tokens)
    return gpu_seconds(n, max_new_tokens)


def text_of(content):
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        return content.get("text") or ""
    if isinstance(content, (list, tuple)):
        return "".join(text_of(c) for c in content)
    return ""


def build_messages(message, history, system_prompt):
    """Chat history -> template messages; collapsed thinking becomes reasoning_content."""
    msgs = []
    reasoning = None
    for turn in history or []:
        role = turn.get("role")
        text = text_of(turn.get("content"))
        if role == "assistant" and (turn.get("metadata") or {}).get("title"):
            reasoning = text
            continue
        if role == "user":
            msgs.append({"role": "user", "content": text})
            reasoning = None
        elif role == "assistant":
            if msgs and msgs[-1]["role"] == "assistant":
                msgs[-1]["content"] += "\n\n" + text
            else:
                m = {"role": "assistant", "content": text}
                if reasoning:
                    m["reasoning_content"] = reasoning
                msgs.append(m)
            reasoning = None
    msgs.append({"role": "user", "content": text_of(message)})
    system = (system_prompt or "").strip()
    return ([{"role": "system", "content": system}] if system else []), msgs


def encode(system, msgs, max_new_tokens):
    limit = MAX_CONTEXT - int(max_new_tokens)
    while True:
        ids = tokenizer.apply_chat_template(
            system + msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True
        )["input_ids"][0]
        if ids.shape[0] <= limit or len(msgs) <= 1:
            break
        msgs = msgs[2:] if len(msgs) > 2 else msgs[1:]  # drop the oldest exchange
    if ids.shape[0] > limit:
        raise gr.Error(
            f"The message is too long for this demo: {MAX_CONTEXT:,} tokens are shared by the prompt and the reply."
        )
    return ids


def split_thinking(text):
    """Flint decides per turn whether to reason; reasoning comes first, inside <think>...</think>."""
    s = text.lstrip()
    if not s.startswith("<think>"):
        if "<think>".startswith(s):  # still undecided
            return "", "", False
        return "", text.strip(), True
    if "</think>" in s:
        thinking, answer = s[len("<think>"):].split("</think>", 1)
        return thinking.strip(), answer.strip(), True
    return s[len("<think>"):].strip(), "", False


def render(thinking, answer, done_thinking, think_secs):
    out = []
    if thinking:
        title = f"Thought for {think_secs:.1f} s" if done_thinking else "Thinking"
        out.append(
            gr.ChatMessage(
                role="assistant",
                content=thinking,
                metadata={"title": title, "status": "done" if done_thinking else "pending"},
            )
        )
    if answer or not out:
        out.append(gr.ChatMessage(role="assistant", content=answer))
    return out


@spaces.GPU(size="xlarge", duration=gpu_duration)
def chat(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k):
    if not text_of(message).strip():
        raise gr.Error("Type a message first.")
    t0 = time.time()
    max_new_tokens = int(max_new_tokens)
    system, msgs = build_messages(message, history, system_prompt)
    input_ids = encode(system, msgs, max_new_tokens)

    budget = gpu_seconds(int(input_ids.shape[0]), max_new_tokens) - 5
    tokens = queue.SimpleQueue()  # token ids from the generation thread; None marks the end
    state = {"n": 0, "error": None, "timed_out": False}

    def should_stop():
        if time.time() - t0 > budget:
            state["timed_out"] = True
            return True
        return False

    def worker():
        try:
            state["n"] = generate(
                model,
                input_ids,
                max_new_tokens=max_new_tokens,
                eos_ids=EOS_IDS,
                temperature=float(temperature),
                top_p=float(top_p),
                top_k=int(top_k),
                on_token=tokens.put,
                should_stop=should_stop,
            )
        except Exception as e:  # surfaced in the UI below
            state["error"] = e
        finally:
            tokens.put(None)

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()

    # Decode in this thread, and at most ~20 UI updates a second, so the GPU loop never waits.
    ids = []
    text = ""
    think_start = None
    think_secs = 0.0
    finished = False
    while not finished:
        try:
            item = tokens.get(timeout=0.05)
        except queue.Empty:
            continue
        if item is None:
            finished = True
        else:
            ids.append(item)
        deadline = time.time() + 0.05
        while not finished and time.time() < deadline:
            try:
                item = tokens.get(timeout=max(0.0, deadline - time.time()))
            except queue.Empty:
                break
            if item is None:
                finished = True
            else:
                ids.append(item)
        if not ids:
            continue
        text = tokenizer.decode(ids, skip_special_tokens=False)
        if text.endswith("\ufffd") and not finished:  # half of a multi-byte character
            continue
        thinking, answer, done = split_thinking(text)
        if thinking and not done:
            think_start = think_start or time.time()
            think_secs = time.time() - think_start
        yield render(thinking, answer, done, think_secs), ""
    thread.join()
    if state["error"] is not None:
        raise gr.Error("Generation failed. Please try again.")

    dt = time.time() - t0
    thinking, answer, _ = split_thinking(text)
    n = state["n"]
    took = f"Took {dt:.1f} s · {n} tokens · {n / dt:.1f} tok/s"
    if state["timed_out"]:
        took += " · stopped at the time limit"
    elif n >= max_new_tokens:
        took += " · hit the max new tokens limit"
    yield render(thinking, answer, True, think_secs), took


EXAMPLES = [
    "Write a Python function that merges overlapping intervals, with a few tests.",
    "Explain what this does and find the bug:\n\n```python\ndef dedupe(items):\n    seen = set()\n    for i, x in enumerate(items):\n        if x in seen:\n            items.pop(i)\n        seen.add(x)\n    return items\n```",
    "Implement an LRU cache in TypeScript with O(1) get and put.",
    "Write a bash one-liner that finds the 10 largest files under the current directory.",
    "Convert this SQL to a pandas expression: SELECT city, AVG(price) FROM sales WHERE year = 2024 GROUP BY city ORDER BY 2 DESC;",
]

with gr.Blocks(title=TITLE) as demo:
    mfn_hf.header(
        TITLE,
        model=MODEL_ID,
        description="Recursive's 40B mixture-of-experts coding model (about 2.5B active per token). It decides on each turn whether to reason first.",
        flow_fn="/chat",
    )

    system_prompt = gr.Textbox(label="System prompt", lines=2, value=DEFAULT_SYSTEM, placeholder="Optional", render=False)
    max_new_tokens = gr.Slider(256, MAX_NEW_TOKENS, value=4096, step=256, label="Max new tokens",
                               info="Reasoning counts toward this.", render=False)
    temperature = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature", info="0 means greedy.", render=False)
    top_p = gr.Slider(0.05, 1.0, value=0.95, step=0.05, label="Top-p", render=False)
    top_k = gr.Slider(0, 200, value=20, step=1, label="Top-k", info="0 turns it off.", render=False)
    took = gr.Markdown("", elem_classes="mfn-meta", render=False)

    gr.ChatInterface(
        chat,
        additional_inputs=[system_prompt, max_new_tokens, temperature, top_p, top_k],
        additional_inputs_accordion=gr.Accordion("Settings", open=False, render=False),
        additional_outputs=[took],
        examples=[[e] for e in EXAMPLES],
        cache_examples=False,
        chatbot=gr.Chatbot(height=560, label="Chat"),
        textbox=gr.Textbox(placeholder="Ask a coding question", show_label=False),
        api_name="chat",
        fill_height=False,
    )
    took.render()
    gr.Markdown("Model by [Recursive](https://www.recursive.com) · Apache-2.0", elem_classes="mfn-meta")
    mfn_hf.footer()


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