#!/usr/bin/env python3
# =============================================================================
# chat_jirack.py - chat with JiRackDeltaNet 70B from a safetensors folder.
#
#   python3 chat_jirack.py --model-dir /mnt/nfs_share/JiRackDeltaNet_70b/expanded
#
# - model code: JiRackDeltaNet_*.py from the same folder (or --model-file)
# - several GPUs: the TransformerBlocks are split across them evenly (70B bf16 = 140 GB,
#   so at least 2x80 GB); 1 GPU if it fits; CPU works but is VERY slow for 70B
# - answer streams token by token, stops on <|im_end|> / <|endoftext|>
# - commands: /reset (clear history), /think (thinking on/off), exit
# =============================================================================
import argparse
import glob
import importlib.util
import itertools
import json
import os
import sys
import time

import torch
import torch.nn as nn

INIT_FNS = ("uniform_", "normal_", "trunc_normal_", "kaiming_uniform_", "kaiming_normal_",
            "xavier_uniform_", "xavier_normal_", "orthogonal_")


def find_model_file(d):
    files = sorted(glob.glob(os.path.join(d, "JiRackDeltaNet_*.py")))
    if not files:
        sys.exit(f"❌ no JiRackDeltaNet_*.py in {d}; pass --model-file")
    return files[-1]


def load_module(path):
    spec = importlib.util.spec_from_file_location("jirack_chat_model", path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules["jirack_chat_model"] = mod
    spec.loader.exec_module(mod)
    return mod


def build(mod):
    """bf16 skeleton without random init (all weights are loaded afterwards) - as in deltanet_gate.py."""
    saved = {n: getattr(nn.init, n) for n in INIT_FNS if hasattr(nn.init, n)}
    for n in saved:
        setattr(nn.init, n, lambda t, *a, **k: t)
    torch.set_default_dtype(torch.bfloat16)
    try:
        model = mod.JiRackQwen38ForCausalLM()
    finally:
        torch.set_default_dtype(torch.float32)
        for n, f in saved.items():
            setattr(nn.init, n, f)
    return model.eval()


def load_weights(model, d):
    from safetensors import safe_open
    idx = os.path.join(d, "model.safetensors.index.json")
    files = sorted(set(json.load(open(idx))["weight_map"].values())) if os.path.exists(idx) else \
        sorted(os.path.basename(p) for p in glob.glob(os.path.join(d, "*.safetensors")))
    sd = model.state_dict()
    params = {n for n, _ in model.named_parameters()}
    loaded = set()
    t0 = time.time()
    with torch.no_grad():
        for i, f in enumerate(files):
            with safe_open(os.path.join(d, f), framework="pt", device="cpu") as h:
                for k in h.keys():
                    if k in sd:
                        sd[k].copy_(h.get_tensor(k))
                        loaded.add(k)
            print(f"   {i + 1}/{len(files)} {f}  {time.time() - t0:.0f} s", flush=True)
    missing = sorted(params - loaded)
    if missing:
        sys.exit(f"❌ {len(missing)} parameters not in {d}: {missing[:5]}")


class ChatTokenizer:
    """tokenizer.json (tokenizers) + chat template (jinja2) - no transformers needed, so the
    tokenizer_config.json format (e.g. extra_special_tokens as a list) does not matter."""

    def __init__(self, d):
        import jinja2
        from jinja2.sandbox import ImmutableSandboxedEnvironment
        from tokenizers import Tokenizer
        self.tk = Tokenizer.from_file(os.path.join(d, "tokenizer.json"))
        cfg = {}
        if os.path.exists(os.path.join(d, "tokenizer_config.json")):
            cfg = json.load(open(os.path.join(d, "tokenizer_config.json"), encoding="utf-8"))
        path = os.path.join(d, "chat_template.jinja")
        text = open(path, encoding="utf-8").read() if os.path.exists(path) else cfg.get("chat_template")
        if isinstance(text, list):
            text = next(t["template"] for t in text if t.get("name") == "default")
        if not text:
            sys.exit(f"❌ no chat template in {d}")

        def raise_exception(msg):
            raise jinja2.exceptions.TemplateError(msg)

        def tojson(x, ensure_ascii=False, indent=None, separators=None, sort_keys=False):
            return json.dumps(x, ensure_ascii=ensure_ascii, indent=indent, separators=separators,
                              sort_keys=sort_keys)

        env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)
        env.filters["tojson"] = tojson
        env.globals["raise_exception"] = raise_exception
        self.tmpl = env.from_string(text)
        s = lambda v: v.get("content") if isinstance(v, dict) else v
        self.bos, self.eos = s(cfg.get("bos_token")) or "", s(cfg.get("eos_token")) or ""
        self.special = [t.content for t in self.tk.get_added_tokens_decoder().values() if t.special]
        self.stop_ids = {i for i in (self.tk.token_to_id("<|im_end|>"), self.tk.token_to_id("<|endoftext|>"),
                                     self.tk.token_to_id(self.eos) if self.eos else None) if i is not None}

    def prompt_ids(self, msgs, think):
        text = self.tmpl.render(messages=msgs, add_generation_prompt=True, enable_thinking=think,
                                bos_token=self.bos, eos_token=self.eos)
        return self.tk.encode(text, add_special_tokens=False).ids

    def decode(self, ids):
        return self.tk.decode(ids, skip_special_tokens=False)


def _to(x, dev):
    if torch.is_tensor(x):
        return x.to(dev)
    if isinstance(x, (tuple, list)):
        return type(x)(_to(v, dev) for v in x)
    if isinstance(x, dict):
        return {k: _to(v, dev) for k, v in x.items()}
    return x


def place(model, devices):
    """Blocks split evenly over the GPUs (moved there directly - the whole model never sits on one
    GPU); everything else on the first GPU. Pre-hooks move each module's inputs to its device,
    so the model code needs no changes."""
    blocks = [m for m in model.modules() if type(m).__name__ == "TransformerBlock"]
    in_block = {id(x) for b in blocks for x in b.modules()}
    for i, b in enumerate(blocks):
        dev = devices[i * len(devices) // len(blocks)]
        b.to(dev)
        if len(devices) > 1:
            b.register_forward_pre_hook(lambda m, a, k, dev=dev: (_to(a, dev), _to(k, dev)), with_kwargs=True)
    for m in model.modules():                             # embeddings / norms / lm_head / mtp / rotary
        if id(m) in in_block:
            continue
        m._apply(lambda t: t.to(devices[0]), recurse=False)
        if len(devices) > 1 and not list(m.children()):
            m.register_forward_pre_hook(lambda mm, a, k: (_to(a, devices[0]), _to(k, devices[0])),
                                        with_kwargs=True)
    return devices[0]


def _cast(x, dtype):
    if torch.is_tensor(x):
        return x.to(dtype) if x.is_floating_point() and x.dtype != dtype else x
    if isinstance(x, (tuple, list)):
        return type(x)(_cast(v, dtype) for v in x)
    if isinstance(x, dict):
        return {k: _cast(v, dtype) for k, v in x.items()}
    return x


def fp32_blocks(model):
    """CPU (AVX-only Xeon has no fast bf16 matmul): every TransformerBlock and lm_head runs with
    fp32 weights/inputs, converted on entry and back on exit - stored weights stay bf16 (RAM),
    same as the 'fp32blocks' compute of deltanet_gate.py."""
    def pre(mod, args, kwargs):
        saved = []
        for t in itertools.chain(mod.parameters(), mod.buffers()):
            if t.is_floating_point() and t.dtype != torch.float32:
                saved.append((t, t.data))
                t.data = t.data.float()
        mod._jr_saved = saved
        return _cast(args, torch.float32), _cast(kwargs, torch.float32)

    def post(mod, args, out, back=torch.bfloat16):
        for t, d in mod._jr_saved:
            t.data = d
        mod._jr_saved = []
        return _cast(out, back) if back is not None else out

    for m in model.modules():
        if type(m).__name__ == "TransformerBlock":
            m.register_forward_pre_hook(pre, with_kwargs=True)
            m.register_forward_hook(post)
    head = getattr(model, "lm_head", None)
    if isinstance(head, nn.Module):
        head.register_forward_pre_hook(pre, with_kwargs=True)
        head.register_forward_hook(lambda m, a, o: post(m, a, o, None))


def logits_of(out):
    if isinstance(out, (tuple, list)):
        return out[0]
    return out.logits if hasattr(out, "logits") else out


def sample(logits, temperature, top_p):
    if temperature <= 0:
        return int(logits.argmax())
    probs = torch.softmax(logits.float() / temperature, dim=-1)
    sp, si = probs.sort(descending=True)
    cut = (sp.cumsum(-1) - sp) > top_p
    sp[cut] = 0
    return int(si[torch.multinomial(sp / sp.sum(), 1)])


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", default="/mnt/nfs_share/JiRackDeltaNet_70b/expanded")
    p.add_argument("--model-file", default=None)
    p.add_argument("--tokenizer-dir", default=None, help="default: --model-dir")
    p.add_argument("--max-new", type=int, default=2048)
    p.add_argument("--temperature", type=float, default=0.6, help="0 = greedy")
    p.add_argument("--top-p", type=float, default=0.95)
    p.add_argument("--system", default=None, help="file with a system prompt")
    p.add_argument("--no-think", action="store_true", help="start with thinking off")
    p.add_argument("--cpu", action="store_true")
    p.add_argument("--cpu-bf16", action="store_true", help="CPU: plain bf16 instead of fp32 blocks")
    a = p.parse_args()

    tok = ChatTokenizer(a.tokenizer_dir or a.model_dir)
    stop = tok.stop_ids
    print(f"stop tokens {sorted(stop)}", flush=True)

    mf = a.model_file or find_model_file(a.model_dir)
    print(f"model file {mf}\nweights    {a.model_dir}", flush=True)
    model = build(load_module(mf))
    load_weights(model, a.model_dir)
    n = sum(q.numel() for q in model.parameters())
    if a.cpu or not torch.cuda.is_available():
        dev = torch.device("cpu")
        try:
            torch.backends.nnpack.set_flags(False)       # silences "Could not initialize NNPACK"
        except Exception:
            pass
        if not a.cpu_bf16:
            fp32_blocks(model)
        print(f"⚠️  CPU: {n / 1e9:.1f}B params, compute {'bf16' if a.cpu_bf16 else 'fp32 blocks'} - "
              f"every token is a full forward pass, expect minutes per token. Ctrl+C stops the answer.")
    else:
        devs = [torch.device("cuda", i) for i in range(torch.cuda.device_count())]
        dev = place(model, devs)
        print(f"on {len(devs)} GPU(s): {n / 1e9:.1f}B params, "
              f"{[round(torch.cuda.memory_allocated(d) / 2 ** 30) for d in devs]} GB")

    system = open(a.system, encoding="utf-8").read().strip() if a.system else None
    think = not a.no_think
    history = []
    print("\nReady. /reset - clear history, /think - thinking on/off, exit - quit")
    while True:
        try:
            user = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not user:
            continue
        if user.lower() in ("exit", "quit"):
            break
        if user == "/reset":
            history = []
            print("(history cleared)")
            continue
        if user == "/think":
            think = not think
            print(f"(thinking {'on' if think else 'off'})")
            continue
        history.append({"role": "user", "content": user})
        msgs = ([{"role": "system", "content": system}] if system else []) + history
        ids = torch.tensor([tok.prompt_ids(msgs, think)], dtype=torch.long, device=dev)
        start = ids.shape[1]
        print("\nAssistant: ", end="", flush=True)
        printed, t0 = "", time.time()
        try:
            with torch.no_grad():
                for _ in range(a.max_new):
                    nxt = sample(logits_of(model(ids))[0, -1], a.temperature, a.top_p)
                    if nxt in stop:
                        break
                    ids = torch.cat([ids, torch.tensor([[nxt]], device=ids.device)], dim=1)
                    text = tok.decode(ids[0, start:].tolist())
                    if not text.endswith("�"):           # wait for a complete UTF-8 character
                        print(text[len(printed):], end="", flush=True)
                        printed = text
        except KeyboardInterrupt:
            print(" [stopped]", end="")
        gen = ids.shape[1] - start
        raw = tok.decode(ids[0, start:].tolist())
        print(f"\n   [{gen} tokens, {gen / max(time.time() - t0, 1e-6):.2f} tok/s]")
        # Qwen: earlier turns go back into the history without the thinking part
        reply = raw.split("</think>")[-1]
        for t in tok.special:
            reply = reply.replace(t, "")
        history.append({"role": "assistant", "content": reply.strip()})


if __name__ == "__main__":
    main()