"""Build the files a faster backend loads, from a model directory (only read) into a new directory.

    python -m decisionnode_edge.build openvino path/to/model --out path/to/model-openvino [--lengths 512,2048]
    python -m decisionnode_edge.build gguf path/to/model --out path/to/model-gguf --llama-cpp path/to/llama.cpp

openvino  The backbone's text path (the token embedding, the layers, the final norm: input_ids [1, L] to the hidden
          states [1, L, hidden]) exported with torch.onnx's dynamo exporter in float32 at each static length, then
          saved as OpenVINO IR with float32 weights (the bf16 weights, exactly); the runtime runs it in bf16. One IR per
          length, since the export unrolls the delta rule over the length: each takes 15 to 20 minutes, and lengths
          above a few thousand tokens take hours. Then each IR is run once against torch. Needs the `openvino` extra
          (openvino, onnx, onnxscript); `--device cuda` (the default when a GPU is present) keeps the weights on the card.
gguf      llama.cpp's converter (convert_hf_to_gguf.py --no-mtp) on links to the backbone files, Q8_0 by default: 8-bit
          keeps the answers, 4-bit does not. Needs a llama.cpp checkout and its converter's Python requirements, in the
          environment of `--python` (default: this one).

Both write decisionnode_edge_build.json (the backend, the files, the model's fingerprint, versions, times), which
EdgeModel.load reads from `artifacts`.
"""

from __future__ import annotations

import argparse
import importlib.metadata
import json
import os
import shutil
import subprocess
import sys
import time

from .backends import BUILD_FILE, BUILD_FORMAT, fingerprint, link_model_files

DEFAULT_LENGTHS = (512, 2048)


def versions(*names: str) -> dict:
    out = {}
    for name in names:
        try:
            out[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            out[name] = None
    return out


def new_directory(out: str) -> None:
    if os.path.exists(out):
        raise SystemExit(f"{out} exists: the build writes a new directory")
    os.makedirs(out)


def write_record(out: str, record: dict) -> None:
    with open(os.path.join(out, BUILD_FILE), "w", encoding="utf-8") as f:
        json.dump(record, f, indent=2)
        f.write("\n")


def model_name(model: str) -> str:
    with open(os.path.join(model, "decisionnode_edge.json"), encoding="utf-8") as f:
        return json.load(f)["name"]


def build_openvino(model: str, out: str, lengths: list[int], device: str | None) -> dict:
    import numpy as np
    import openvino as ov
    import torch
    from transformers import AutoModelForImageTextToText, AutoTokenizer

    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    record = {"format": BUILD_FORMAT, "backend": "openvino", "model": model_name(model), "model_fingerprint": fingerprint(model), "files": {},
              "weights": "float32 (the bf16 weights, exactly); run in bf16", "export_device": device,
              "versions": versions("torch", "transformers", "onnx", "onnxscript", "openvino"), "seconds": {}, "check": {}}  # fmt: skip
    tokenizer = AutoTokenizer.from_pretrained(model)
    record["pad_token_id"] = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

    class LastHidden(torch.nn.Module):
        def __init__(self, text):
            super().__init__()
            self.text = text

        def forward(self, input_ids):
            return self.text(input_ids=input_ids, use_cache=False).last_hidden_state

    full = AutoModelForImageTextToText.from_pretrained(model, dtype=torch.float32).to(device).eval()
    text = LastHidden(full.model.language_model).eval()
    generator = torch.Generator().manual_seed(0)
    for length in lengths:
        work = os.path.join(out, f"onnx-{length}")
        os.makedirs(work)
        onnx_file = os.path.join(work, "edge.onnx")
        ids = torch.randint(10, 200000, (1, length), generator=generator).to(device)
        t0 = time.perf_counter()
        with torch.inference_mode():
            torch.onnx.export(text, (ids,), onnx_file, dynamo=True, external_data=True, input_names=["input_ids"], output_names=["last_hidden_state"], optimize=True)
        t1 = time.perf_counter()
        ir = ov.convert_model(onnx_file)
        name = f"openvino-{length}.xml"
        ov.save_model(ir, os.path.join(out, name), compress_to_fp16=False)
        del ir
        shutil.rmtree(work)
        record["files"][str(length)] = name
        record["seconds"][str(length)] = {"export": round(t1 - t0, 1), "convert": round(time.perf_counter() - t1, 1)}
        print(json.dumps({"length": length, **record["seconds"][str(length)]}), flush=True)

    # each IR against torch (float32) on a prompt shorter than its length, so the padding is read past as at run time
    core = ov.Core()
    for length in lengths:
        n = length - 37
        ids = torch.randint(10, 200000, (1, n), generator=generator)
        with torch.inference_mode():
            want = text(ids.to(device))[0, -1].float().cpu()
        compiled = core.compile_model(os.path.join(out, record["files"][str(length)]), "CPU", {"INFERENCE_PRECISION_HINT": "bf16"})
        x = np.full((1, length), record["pad_token_id"], dtype=np.int64)
        x[0, :n] = ids[0].numpy()
        got = torch.from_numpy(np.asarray(compiled.create_infer_request().infer({0: x})[compiled.output(0)][0, n - 1], dtype=np.float32))
        cosine = float(torch.nn.functional.cosine_similarity(got, want, dim=0))
        record["check"][str(length)] = {"tokens": n, "cosine_to_torch_float32": round(cosine, 6), "largest_difference": round(float((got - want).abs().max()), 4)}
        print(json.dumps({"check": length, **record["check"][str(length)]}), flush=True)
        del compiled
        if not torch.isfinite(got).all() or cosine < 0.99:
            raise SystemExit(f"the IR at {length} tokens does not reproduce the backbone (cosine {cosine:.4f}): {out} is incomplete")
    return record


def build_gguf(model: str, out: str, llama_cpp: str, outtype: str, python: str) -> dict:
    converter = os.path.join(llama_cpp, "convert_hf_to_gguf.py")
    if not os.path.isfile(converter):
        raise SystemExit(f"{converter} not found: --llama-cpp is the llama.cpp checkout")
    source = os.path.join(out, "source")  # the converter reads links to the backbone files, never the readout
    os.makedirs(source)
    link_model_files(model, source)
    name = f"model-{outtype}.gguf"
    t0 = time.perf_counter()
    subprocess.run([python, converter, source, "--no-mtp", "--outfile", os.path.join(out, name), "--outtype", outtype], check=True)
    shutil.rmtree(source)
    commit = subprocess.run(["git", "-C", llama_cpp, "rev-parse", "--short", "HEAD"], capture_output=True, text=True).stdout.strip() or None
    return {"format": BUILD_FORMAT, "backend": "llamacpp", "model": model_name(model), "model_fingerprint": fingerprint(model), "files": {"gguf": name},
            "type": outtype, "bytes": os.path.getsize(os.path.join(out, name)), "llama_cpp": commit, "seconds": round(time.perf_counter() - t0, 1)}  # fmt: skip


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="python -m decisionnode_edge.build", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("backend", choices=("openvino", "gguf"))
    parser.add_argument("model", help="the model directory (only read)")
    parser.add_argument("--out", required=True, help="the directory to write (must not exist)")
    parser.add_argument("--lengths", default=",".join(map(str, DEFAULT_LENGTHS)), help="openvino: the static lengths in tokens (default 512,2048)")
    parser.add_argument("--device", default=None, help="openvino: where the export runs, cuda or cpu (default: cuda when present)")
    parser.add_argument("--llama-cpp", help="gguf: the llama.cpp checkout (its convert_hf_to_gguf.py)")
    parser.add_argument("--type", default="q8_0", choices=("q8_0", "bf16"), help="gguf: the weight type (default q8_0)")
    parser.add_argument("--python", default=sys.executable, help="gguf: the Python that runs llama.cpp's converter (default: this one)")
    args = parser.parse_args(argv)
    for name in ("config.json", "model.safetensors", "decisionnode_edge.json"):
        if not os.path.isfile(os.path.join(args.model, name)):
            raise SystemExit(f"{args.model} is not a DecisionNode Edge model directory (no {name})")
    if args.backend == "gguf" and not args.llama_cpp:
        raise SystemExit("gguf needs --llama-cpp")
    new_directory(args.out)
    if args.backend == "openvino":
        lengths = sorted({int(x) for x in args.lengths.split(",") if x.strip()})
        record = build_openvino(args.model, args.out, lengths, args.device)
    else:
        record = build_gguf(args.model, args.out, args.llama_cpp, args.type, args.python)
    write_record(args.out, record)
    print(json.dumps({k: v for k, v in record.items() if k != "versions"}), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())