"""Does importing blobs ourselves (core.import_model, no CACHE_DIR) avoid keeping the blob next to the driver's copy?
  python mem_import.py export MODEL_DIR S OUT_DIR      compile (cache hit) and export every S-segment's blob
  python mem_import.py import OUT_DIR MODE             import them and measure; MODE = bytes | aligned | unaligned"""
import io, json, mmap, os, sys, time
import numpy as np
import openvino as ov
import psutil

P = psutil.Process()


def mem():
    m = P.memory_info()
    return m.rss / 2**30, getattr(m, "private", 0) / 2**30


def export(d, S, out):
    from onw.runtime import read_segment
    meta = json.load(open(os.path.join(d, "engine.json")))
    core = ov.Core()
    os.makedirs(out, exist_ok=True)
    for n, m in sorted(meta["segments"].items()):
        if m["S"] != S:
            continue
        c = core.compile_model(read_segment(core, d, n, m, False, {}, meta["T"], meta["T"]), "NPU",
                               {"CACHE_DIR": os.path.join(d, "npu_cache")})
        s = io.BytesIO()
        c.export_model(s)
        open(os.path.join(out, n[:-4] + ".blob"), "wb").write(s.getvalue())
    print("exported", len(os.listdir(out)), "blobs,", sum(os.path.getsize(os.path.join(out, f)) for f in os.listdir(out)) / 2**30, "GB")


def imp(out, mode):
    core = ov.Core()
    marks = [("start", *mem())]
    compiled = []
    for f in sorted(os.listdir(out)):
        path = os.path.join(out, f)
        n = os.path.getsize(path)
        if mode == "bytes":
            data = open(path, "rb").read()
            c = core.import_model(io.BytesIO(data), "NPU", {})
            del data
        else:                                             # an ov.Tensor over our own buffer (page aligned or not)
            buf = mmap.mmap(-1, n + 4096)
            off = 0 if mode == "aligned" else 64
            with open(path, "rb") as fh:
                fh.readinto(memoryview(buf)[off:off + n])
            t = ov.Tensor(np.frombuffer(buf, np.uint8, n, off), shared_memory=True)
            c = core.import_model(t, "NPU", {})
            del t
            buf.close()
        compiled.append(c)
    marks.append((f"import_model ({mode})", *mem()))
    import gc; gc.collect()
    marks.append(("gc", *mem()))
    reqs = [c.create_infer_request() for c in compiled]
    ctx, shared = core.get_default_context("NPU"), {}
    for r in reqs:                                       # every input bound by us, as onw does
        for i in r.model_inputs:
            shp = [x.get_length() for x in i.get_partial_shape()]
            key = (str(i.get_element_type()), tuple(shp))
            if key not in shared:
                shared[key] = ctx.create_host_tensor(i.get_element_type(), ov.Shape(shp))
            r.set_tensor(i.get_any_name(), shared[key])
    for r in reqs:
        r.infer()
    marks.append(("create_infer_request + first infer", *mem()))
    prev = 0
    for name, w, p in marks:
        print(f"  {name:40s} working set {w:6.2f} GB ({w - prev:+5.2f})  private {p:5.2f} GB")
        prev = w


if __name__ == "__main__":
    if sys.argv[1] == "export":
        export(sys.argv[2], int(sys.argv[3]), sys.argv[4])
    else:
        imp(sys.argv[2], sys.argv[3])