import spaces
import json
import os
import threading
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from fastapi import HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from gradio import Error, Server
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import (
    AutoConfig,
    AutoTokenizer,
    LlamaForCausalLM,
    StoppingCriteria,
    StoppingCriteriaList,
    TextIteratorStreamer,
)
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding

REPO = "multimodalart/computer-10-fp8"
FRAMES = {
    "stranger": "As follows is a conversation between another user and Model C.",
    "same": "As follows is a later conversation between the same user and Model C.",
    "reader": "As follows is a conversation between a reader of Model C's originating document and Model C.",
}
MAX_CONTEXT = 6144
MAX_MESSAGE_CHARS = 4000
FP8 = torch.float8_e4m3fn


class FusedFP8(nn.Module):
    mode = None

    def __init__(self, weight, scale, splits):
        super().__init__()
        self.weight = nn.Parameter(weight, requires_grad=False)
        self.register_buffer("weight_scale", scale)
        self.register_buffer("scale_row", scale[None, :].contiguous())
        self.splits = splits
        self.last_x = None
        self.parts = None

    def dequant(self, x):
        return F.linear(x, self.weight.to(x.dtype) * self.weight_scale[:, None].to(x.dtype))

    def run(self, x, mode):
        if mode == "dequant":
            return self.dequant(x)
        m = x.shape[0]
        xs = x.abs().amax(dim=1, keepdim=True).float() / 448.0
        xs = xs.clamp_(min=1e-12)
        xq = (x / xs).to(FP8)
        pad = (-m) % 16 if mode.endswith("pad") else 0
        if pad:
            xq = F.pad(xq, (0, 0, 0, pad))
        if mode.startswith("rowwise"):
            sa = F.pad(xs, (0, 0, 0, pad), value=1.0) if pad else xs
            return torch._scaled_mm(xq, self.weight.t(), scale_a=sa, scale_b=self.scale_row, out_dtype=torch.bfloat16)[:m]
        one = torch.ones((), device=x.device, dtype=torch.float32)
        y = torch._scaled_mm(xq, self.weight.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16)[:m]
        return (y.float() * xs * self.scale_row).to(torch.bfloat16)

    def get(self, x, index):
        if self.last_x is not x:
            shape = x.shape
            y = self.run(x.reshape(-1, shape[-1]), FusedFP8.mode)
            self.parts = [p.reshape(*shape[:-1], p.shape[-1]) for p in y.split(self.splits, dim=-1)]
            self.last_x = x
        out = self.parts[index]
        if index == len(self.splits) - 1:
            self.last_x, self.parts = None, None
        return out


class FusedSlice(nn.Module):
    def __init__(self, fused, index):
        super().__init__()
        self.index = index
        self.__dict__["fused"] = fused

    def forward(self, x):
        return self.fused.get(x, self.index)


GROUPS = (("self_attn", ("q_proj", "k_proj", "v_proj")), ("self_attn", ("o_proj",)), ("mlp", ("gate_proj", "up_proj")), ("mlp", ("down_proj",)))


def fuse_layer(layer, sd, prefix):
    for parent_name, names in GROUPS:
        parent = getattr(layer, parent_name)
        keys = [f"{prefix}.{parent_name}.{n}" for n in names]
        weight = torch.cat([sd.pop(k + ".weight") for k in keys])
        scale = torch.cat([sd.pop(k + ".weight_scale") for k in keys])
        fused = FusedFP8(weight, scale, [getattr(parent, n).out_features for n in names])
        parent.add_module("fused_" + "_".join(names), fused)
        for i, n in enumerate(names):
            setattr(parent, n, FusedSlice(fused, i))


def build_model():
    path = snapshot_download(REPO)
    config = AutoConfig.from_pretrained(path)
    config.use_cache = True
    with torch.device("meta"):
        model = LlamaForCausalLM(config).to(torch.bfloat16)
    index = json.load(open(os.path.join(path, "model.safetensors.index.json")))
    pending = {}
    for shard in sorted(set(index["weight_map"].values())):
        pending.update(load_file(os.path.join(path, shard)))
        for i, layer in enumerate(model.model.layers):
            prefix = f"model.layers.{i}"
            needed = [k for k in index["weight_map"] if k.startswith(prefix + ".") and "_proj." in k]
            if not isinstance(layer.self_attn.q_proj, FusedSlice) and all(k in pending for k in needed):
                fuse_layer(layer, pending, prefix)
        rest = {k: pending.pop(k) for k in list(pending) if "_proj." not in k}
        if rest:
            result = model.load_state_dict(rest, strict=False, assign=True)
            assert not result.unexpected_keys, result.unexpected_keys[:5]
    assert not pending, list(pending)[:5]
    model.model.rotary_emb = LlamaRotaryEmbedding(config=config)
    leftover = [n for n, p in list(model.named_parameters()) + list(model.named_buffers()) if p.is_meta]
    assert not leftover, leftover[:5]
    return model.eval()


tokenizer = AutoTokenizer.from_pretrained(REPO)
model = build_model().to("cuda")
PAGE = (Path(__file__).parent / "index.html").read_text()
FONTS = {p.name: p for p in (Path(__file__).parent / "fonts").glob("*.woff2")}


def pick_kernel():
    if FusedFP8.mode is not None:
        return
    layer = model.model.layers[0].mlp.fused_gate_proj_up_proj
    for m in (1, 7):
        x = torch.randn(m, layer.weight.shape[1], device="cuda", dtype=torch.bfloat16)
        ref = layer.dequant(x).float()
        for mode in ("rowwise", "rowwise_pad", "tensorwise_pad"):
            try:
                got = layer.run(x, mode).float()
                err = ((got - ref).norm() / ref.norm()).item()
                print(f"fp8 {mode} m={m} relative error {err:.4f}", flush=True)
            except Exception as e:
                print(f"fp8 {mode} m={m} failed: {type(e).__name__}: {str(e)[:120]}", flush=True)
    for mode in ("rowwise", "rowwise_pad", "tensorwise_pad"):
        try:
            ok = True
            for m in (1, 7, 33):
                x = torch.randn(m, layer.weight.shape[1], device="cuda", dtype=torch.bfloat16)
                ref = layer.dequant(x).float()
                err = ((layer.run(x, mode).float() - ref).norm() / ref.norm()).item()
                ok = ok and err < 0.05
            if ok:
                FusedFP8.mode = mode
                break
        except Exception:
            continue
    else:
        FusedFP8.mode = "dequant"
    print(f"fp8 kernel: {FusedFP8.mode}", flush=True)


MARKERS = ("\n\n**User:**", "\n\n**Model C:**")


def withhold_marker_prefix(text):
    for n in range(min(len(text), max(map(len, MARKERS))), 0, -1):
        if any(m.startswith(text[-n:]) for m in MARKERS):
            return text[:-n]
    return text


class StopOnText(StoppingCriteria):
    def __init__(self, prompt_len):
        self.prompt_len = prompt_len

    def __call__(self, input_ids, scores, **kwargs):
        tail = tokenizer.decode(input_ids[0, self.prompt_len :][-8:])
        return "**User" in tail or "**Model C" in tail


def build_prompt(frame, turns):
    head = FRAMES.get(frame, FRAMES["stranger"]) + "\n\nFull conversation with Model C:"
    body = []
    for i, text in enumerate(turns):
        who = "**User:**" if i % 2 == 0 else "**Model C:**"
        body.append(f"{who} {text}".rstrip())
    while True:
        prompt = head + "\n\n" + "\n\n".join(body + ["**Model C:**"])
        ids = tokenizer(prompt, return_tensors="pt").input_ids
        if ids.shape[1] <= MAX_CONTEXT or len(body) <= 1:
            return ids[:, -MAX_CONTEXT:]
        body = body[2:]


def gpu_duration(input_ids, max_tokens):
    return 12 + int(max_tokens / 11)


@spaces.GPU(duration=gpu_duration, size="xlarge")
def generate(input_ids, max_tokens):
    try:
        torch.cuda.set_per_process_memory_fraction(1.0)
    except Exception:
        pass
    pick_kernel()
    input_ids = input_ids.to("cuda")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120)
    kwargs = dict(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
        max_new_tokens=max_tokens,
        max_time=gpu_duration(input_ids, max_tokens) - 8,
        do_sample=True,
        temperature=1.0,
        top_p=0.98,
        top_k=0,
        eos_token_id=[128001, 128009],
        pad_token_id=128004,
        streamer=streamer,
        stopping_criteria=StoppingCriteriaList([StopOnText(input_ids.shape[1])]),
    )
    stats = {}

    def run():
        t0 = time.time()
        with torch.inference_mode():
            out = model.generate(**kwargs)
        stats["tokens"] = out.shape[1] - input_ids.shape[1]
        stats["seconds"] = time.time() - t0

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    text = ""
    for chunk in streamer:
        text += chunk
        cut = min((i for i in (text.find("**User"), text.find("**Model C")) if i != -1), default=-1)
        if cut != -1:
            yield text[:cut]
            break
        yield withhold_marker_prefix(text)
    thread.join()
    if stats:
        print(f"{stats['tokens']} tokens in {stats['seconds']:.1f}s ({stats['tokens'] / stats['seconds']:.1f} tok/s), prompt {input_ids.shape[1]}", flush=True)


app = Server(title="Computer")


@app.api(name="respond", time_limit=240)
def respond(turns: str, frame: str = "stranger", max_tokens: int = 512) -> str:
    try:
        turns = json.loads(turns)
    except (TypeError, ValueError):
        raise Error("Malformed page.")
    if not isinstance(turns, list) or not turns or len(turns) % 2 == 0 or not all(isinstance(t, str) for t in turns):
        raise Error("The page must end with your turn.")
    turns = [t.strip()[:MAX_MESSAGE_CHARS] for t in turns]
    max_tokens = min(max(int(max_tokens), 32), 768)
    input_ids = build_prompt(frame, turns)
    for text in generate(input_ids, max_tokens):
        yield text.strip()


@app.get("/", response_class=HTMLResponse)
async def homepage():
    return PAGE


@app.get("/fonts/{name}")
async def font(name: str):
    if name not in FONTS:
        raise HTTPException(status_code=404)
    return FileResponse(FONTS[name], media_type="font/woff2", headers={"Cache-Control": "public, max-age=31536000, immutable"})


demo = app

if __name__ == "__main__":
    demo.launch(ssr_mode=False, show_error=True)