#!/usr/bin/env python3
import argparse
import gc
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast


def add_path(path):
    path = str(Path(path).resolve())
    if path not in sys.path:
        sys.path.insert(0, path)


def link_raven(model_dir, raven_dir):
    raven_pkg = raven_dir / "raven"
    if not raven_pkg.is_dir():
        raise FileNotFoundError(f"missing {raven_pkg}")

    for target_dir in [model_dir, *Path.home().glob(".cache/huggingface/modules/transformers_modules/Quasar-Preview*")]:
        if not (target_dir / "modeling_quasar_long.py").exists() and target_dir != model_dir:
            continue
        link = target_dir / "raven"
        if link.is_symlink() or not link.exists():
            link.unlink(missing_ok=True)
            link.symlink_to(raven_pkg, target_is_directory=True)


def load_tokenizer(model_dir):
    # AutoTokenizer fails for this checkpoint because tokenizer_config.json names
    # TokenizersBackend. tokenizer.json itself is valid, so load it directly.
    return PreTrainedTokenizerFast(
        tokenizer_file=str(model_dir / "tokenizer.json"),
        bos_token="<|startoftext|>",
        eos_token="<|endoftext|>",
        pad_token="<|endoftext|>",
        cls_token="[CLS]",
    )


def set_param(root, name, tensor):
    obj = root
    parts = name.split(".")
    for part in parts[:-1]:
        obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
    old = getattr(obj, parts[-1])
    setattr(obj, parts[-1], torch.nn.Parameter(tensor, requires_grad=getattr(old, "requires_grad", False)))


def fill_meta_params(model, model_dir):
    index = json.loads((model_dir / "model.safetensors.index.json").read_text())["weight_map"]
    names = [n for n, p in model.named_parameters() if getattr(p, "is_meta", False)]
    by_file = {}
    for name in names:
        by_file.setdefault(index[name], []).append(name)

    for filename, keys in sorted(by_file.items()):
        tensors = load_file(str(model_dir / filename), device="cpu")
        for key in keys:
            set_param(model, key, tensors[key].to(torch.bfloat16))
        del tensors
        gc.collect()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", default="/workspace/quasar_preview_run/Quasar-Preview")
    p.add_argument("--raven-dir", default="/workspace/quasar_preview_run/goombalab-raven")
    p.add_argument("--prompt", default="Mo Salah is")
    p.add_argument("--max-new-tokens", type=int, default=24)
    p.add_argument("--device", default="cuda:0")
    args = p.parse_args()

    model_dir = Path(args.model_dir).resolve()
    raven_dir = Path(args.raven_dir).resolve()
    add_path(model_dir)
    add_path(raven_dir)
    link_raven(model_dir, raven_dir)

    tokenizer = load_tokenizer(model_dir)
    model = AutoModelForCausalLM.from_pretrained(
        str(model_dir),
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=False,
    ).eval()
    fill_meta_params(model, model_dir)
    model.to(args.device)

    enc = tokenizer(args.prompt, return_tensors="pt")
    enc.pop("token_type_ids", None)
    enc = {k: v.to(args.device) for k, v in enc.items()}

    with torch.inference_mode():
        output_ids = model.generate(
            **enc,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    print(tokenizer.decode(output_ids[0], skip_special_tokens=True))


if __name__ == "__main__":
    main()