"""onw AMD test: are a FastFlowLM model's files intact? Checks each file against the hashes in flm's own
model_info.json (sha256 for the big files, git's blob sha1 for the small ones). Standard library only; reads only.

  python3 flm_verify.py gemma4-it:12b                 the copy in onw's model folder
  python3 flm_verify.py gemma4-it:12b /path/to/flm    another FLM_MODEL_PATH (e.g. Lemonade's, ~/.config/flm)
"""
import glob, hashlib, json, os, sys

WIN = os.name == "nt"


def cfg_dir():
    base = os.environ.get("APPDATA") if WIN else os.path.join(os.path.expanduser("~"), ".config")
    return os.path.join(base or os.path.expanduser("~"), "onw-amd")


def flm_dir():
    exe = "flm.exe" if WIN else "flm"
    hits = glob.glob(os.path.join(cfg_dir(), "bin", "flm", "npu", "**", exe), recursive=True)
    for d in [os.path.dirname(h) for h in hits] + ["/opt/fastflowlm/share/flm", r"C:\Program Files\flm"]:
        if os.path.exists(os.path.join(d, "model_info.json")):
            return d
    sys.exit("model_info.json not found (install FastFlowLM from the onw window first)")


def models_root(arg):
    if arg:
        return os.path.expanduser(arg)
    try:
        d = json.load(open(os.path.join(cfg_dir(), "config.json"), encoding="utf8")).get("models_dir")
    except Exception:
        d = None
    return d or os.path.join(cfg_dir(), "flm")


def digest(path, kind, size):
    h = hashlib.sha256() if kind == "sha256" else hashlib.sha1(f"blob {size}\0".encode())
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(1 << 22), b""):
            h.update(b)
    return h.hexdigest()


def main():
    args = [a for a in sys.argv[1:] if a != "-"]
    ck = args[0] if args else "gemma4-it:12b"
    d = flm_dir()
    family, tag = ck.split(":", 1)
    entry = json.load(open(os.path.join(d, "model_list.json"), encoding="utf8"))["models"][family][tag]
    info = {f["path"]: f for f in json.load(open(os.path.join(d, "model_info.json"), encoding="utf8"))[ck]}
    folder = os.path.join(models_root(args[1] if len(args) > 1 else None), "models", entry["name"])
    print("model", ck, "| folder", folder)
    bad = 0
    for name in entry["files"]:
        f, p = info.get(name), os.path.join(folder, name)
        if not os.path.exists(p):
            print(f"  MISSING  {name}")
            bad += 1
            continue
        size = os.path.getsize(p)
        if f is None:
            print(f"  ?        {name} ({size} bytes; no hash to compare)")
            continue
        want_size = f.get("size")
        kind, want = ("sha256", f["lfs"]["oid"]) if f.get("lfs") else ("sha1", f.get("oid"))
        if want_size and size != want_size:
            print(f"  BAD SIZE {name}: {size} bytes, expected {want_size}")
            bad += 1
            continue
        print(f"  ...      {name} ({size / 2**20:.0f} MB) hashing", end="\r", flush=True)
        got = digest(p, kind, size)
        ok = got == want
        bad += not ok
        print(f"  {'OK      ' if ok else 'BAD HASH'} {name} ({size / 2**20:.0f} MB){'' if ok else f'  {kind} {got[:16]}…, expected {want[:16]}…'}")
    print("\nall files intact" if not bad else f"\n{bad} file(s) damaged: download the model again (onw: delete it in the Models tab, then Download)")


if __name__ == "__main__":
    main()