#!/usr/bin/env python3
"""Read-only scan of the 40250 client packs (batch 2R, docs/PORT-PLAN.md).

Follows the 40250 read chain: `UserInterface.cpp` PackInitialize (pack/Index), `CEterPackManager`
(RegisterPack, the `_texcache` companion, the root pack, `ConvertFileName`), `CEterPack::__BuildIndex`
(.eix: plain `EPKD` or an `MCOZ` TEA+LZO object) and `CEterPack::Get2` (per-entry compressed_type).

  epk_scan.py summary [--json OUT]        per-pack / per-type counts and bytes, registration, overrides, paths
  epk_scan.py verify [--per-type N]       decode N sample entries of every (pack, type) that needs no server key

The index and security keys are read from the reference `EterPack/EterPack.cpp` at run time; nothing
key-like is stored here. Nothing is written except the optional --json report.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import struct
import sys
import zlib
from collections import Counter, defaultdict
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO / ".agents/skills/metin2-40250-parity-audit/scripts"))
import refroot  # noqa: E402

TYPE_NAMES = {0: "NONE", 1: "COMPRESS", 2: "SECURITY", 3: "PANAMA", 4: "HYBRIDCRYPT", 5: "HYBRIDCRYPT_WITHSDB"}
# Types CEterPack::Get2 can decode with keys compiled into the client; 3 needs the IV from the server
# (CAccountConnector RegisterPack(..., abIV)), 4/5 need GC_HYBRIDCRYPT_KEYS / GC_HYBRIDCRYPT_SDB.
LOCAL_TYPES = {0, 1, 2}
EPKD = b"EPKD"
MCOZ = b"MCOZ"
# TEterPackIndex, #pragma pack(4), 32-bit long: id, filename[161], crc, real_data_size, data_size, data_crc,
# data_position, compressed_type. Only data_size (the stored bytes) is read by CEterPack::Get2; real_data_size is
# packer bookkeeping (free-block reuse) and is smaller than data_size in ~half the COMPRESS entries here, so it is
# ignored. The decoded size lives only in the CLZObject header of COMPRESS/SECURITY data.
INDEX = struct.Struct("<i161s3xIiiIib3x")
assert INDEX.size == 192


# ---------------------------------------------------------------- 40250 primitives

def reference_keys(ref: Path) -> dict[str, tuple[int, ...]]:
    text = (ref / "EterPack/EterPack.cpp").read_bytes().decode("latin-1")
    keys = {}
    for name in ("s_adwEterPackKey", "s_adwEterPackSecurityKey"):
        m = re.search(name + r"\[\]\s*=\s*\{([^}]*)\}", text)
        if not m:
            raise SystemExit(f"{name} not found in reference EterPack.cpp")
        keys[name] = tuple(int(v) for v in re.findall(r"\d+", m.group(1)))
    return keys


def tea_decrypt(data: bytes, key: tuple[int, ...]) -> bytes:
    """EterBase/tea.cpp tea_decrypt: 32 rounds, DELTA 0x9E3779B9, 32-bit words, size rounded up to 8."""
    if len(data) % 8:
        data = data + bytes(8 - len(data) % 8)
    words = struct.unpack(f"<{len(data) // 4}I", data)
    out = []
    k0, k1, k2, k3 = key
    kk = (k0, k1, k2, k3)
    M = 0xFFFFFFFF
    for i in range(0, len(words), 2):
        y, z = words[i], words[i + 1]
        s = (0x9E3779B9 * 32) & M
        for _ in range(32):
            z = (z - ((((y << 4) ^ (y >> 5)) + y) ^ (s + kk[(s >> 11) & 3]))) & M
            s = (s - 0x9E3779B9) & M
            y = (y - ((((z << 4) ^ (z >> 5)) + z) ^ (s + kk[s & 3]))) & M
        out += (y, z)
    return struct.pack(f"<{len(out)}I", *out)


def lzo1x_decompress(src: bytes, out_len: int) -> bytes:
    """lzo1x_decompress (lzo1x_d.ch), the decoder 40250 calls in CLZObject::Decompress."""
    op = bytearray()
    ip = 0

    def copy_match(dist_back: int, n: int) -> None:
        start = len(op) - dist_back
        if start < 0:
            raise ValueError("lzo: match before start")
        if dist_back >= n:
            op.extend(op[start:start + n])
        else:
            for j in range(n):
                op.append(op[start + j])

    state = "loop"
    t = 0
    if src[0] > 17:
        t = src[0] - 17
        ip = 1
        if t < 4:
            state = "match_next"
        else:
            op += src[ip:ip + t]
            ip += t
            state = "first_literal_run"
    while True:
        if state == "loop":
            t = src[ip]; ip += 1
            if t >= 16:
                state = "match"
                continue
            if t == 0:
                while src[ip] == 0:
                    t += 255; ip += 1
                t += 15 + src[ip]; ip += 1
            op += src[ip:ip + t + 3]
            ip += t + 3
            state = "first_literal_run"
        if state == "first_literal_run":
            t = src[ip]; ip += 1
            if t >= 16:
                state = "match"
            else:
                dist = 1 + 0x0800 + (t >> 2) + (src[ip] << 2); ip += 1
                copy_match(dist, 3)
                state = "match_done"
        if state == "match":
            if t >= 64:
                dist = 1 + ((t >> 2) & 7) + (src[ip] << 3); ip += 1
                copy_match(dist, (t >> 5) + 1)
            elif t >= 32:
                t &= 31
                if t == 0:
                    while src[ip] == 0:
                        t += 255; ip += 1
                    t += 31 + src[ip]; ip += 1
                dist = 1 + (src[ip] >> 2) + (src[ip + 1] << 6); ip += 2
                copy_match(dist, t + 2)
            elif t >= 16:
                dist = (t & 8) << 11
                t &= 7
                if t == 0:
                    while src[ip] == 0:
                        t += 255; ip += 1
                    t += 7 + src[ip]; ip += 1
                dist += (src[ip] >> 2) + (src[ip + 1] << 6); ip += 2
                if dist == 0:
                    break  # end of stream
                copy_match(dist + 0x4000, t + 2)
            else:
                dist = 1 + (t >> 2) + (src[ip] << 2); ip += 1
                copy_match(dist, 2)
            state = "match_done"
        if state == "match_done":
            t = src[ip - 2] & 3
            if t == 0:
                state = "loop"
                continue
            state = "match_next"
        if state == "match_next":
            op += src[ip:ip + t]
            ip += t
            t = src[ip]; ip += 1
            state = "match"
    if len(op) != out_len:
        raise ValueError(f"lzo: size {len(op)} != {out_len}")
    return bytes(op)


def lz_object(buf: bytes, key: tuple[int, ...] | None) -> bytes:
    """CLZO::Decompress: THeader {fourcc, encrypt, compressed, real}, then MCOZ + LZO (TEA'd when encrypt != 0)."""
    fourcc, enc, comp, real = struct.unpack_from("<4sIII", buf)
    if fourcc != MCOZ:
        raise ValueError("not an MCOZ object")
    if enc:
        if key is None:
            raise ValueError("encrypted object without key")
        body = tea_decrypt(buf[16:16 + enc], key)
        if body[:4] != MCOZ:
            raise ValueError("TEA key incorrect")
        return lzo1x_decompress(body[4:4 + comp], real)
    return lzo1x_decompress(buf[20:20 + comp], real)


# ---------------------------------------------------------------- packs

def read_registration(pack_dir: Path) -> list[tuple[str, str]]:
    """PackInitialize: line 0 is PACK/FILE, then (folder, name) pairs for i in [1, count-1)."""
    lines = (pack_dir / "Index").read_bytes().decode("latin-1").splitlines()
    if lines[0] not in ("PACK", "FILE"):
        raise SystemExit("pack/Index: first line must be PACK or FILE")
    pairs = []
    for i in range(1, len(lines) - 1, 2):
        pairs.append((lines[i], lines[i + 1]))
    return pairs


def read_index(eix: Path, key: tuple[int, ...]) -> tuple[list[dict], bool]:
    raw = eix.read_bytes()
    encrypted = raw[:4] == MCOZ
    data = lz_object(raw, key) if encrypted else raw
    if data[:4] != EPKD:
        raise ValueError(f"{eix.name}: fourcc {data[:4]!r}")
    ver, count = struct.unpack_from("<Ii", data, 4)
    if ver != 2 or len(data) < 12 + INDEX.size * count:
        raise ValueError(f"{eix.name}: version {ver} / size")
    entries = []
    for i in range(count):
        eid, name, crc, real, size, dcrc, pos, ctype = INDEX.unpack_from(data, 12 + INDEX.size * i)
        if not crc:
            continue  # free index slot (PushFreeIndex)
        entries.append({"name": name.split(b"\0", 1)[0].decode("latin-1"), "crc": crc, "block": real,
                        "size": size, "data_crc": dcrc, "pos": pos, "type": ctype})
    return entries, encrypted


def find_file(pack_dir: Path, stem: str, ext: str) -> Path | None:
    """40250 runs on case-insensitive NTFS: `Etc` in Index opens `ETC.eix`."""
    want = (stem + ext).lower()
    for p in pack_dir.iterdir():
        if p.name.lower() == want:
            return p
    return None


def load(client: Path, ref: Path):
    pack_dir = client / "pack"
    keys = reference_keys(ref)
    order = []  # (registered name, folder) in CEterPackManager insertion order
    for folder, name in read_registration(pack_dir):
        order.append((name, folder))
        order.append((name + "_texcache", folder))
    order.append(("root", ""))  # RegisterRootPack after the loop
    packs = {}
    for name, folder in order:
        if name in packs:
            continue  # m_PackMap.find: registered once
        eix, epk = find_file(pack_dir, name, ".eix"), find_file(pack_dir, name, ".epk")
        if not eix or not epk:
            packs[name] = {"folder": folder, "missing": True}
            continue
        entries, encrypted = read_index(eix, keys["s_adwEterPackKey"])
        packs[name] = {"folder": folder, "missing": False, "eix": eix, "epk": epk, "index_encrypted": encrypted,
                       "entries": entries, "epk_bytes": epk.stat().st_size}
    known = {p.name.lower() for n, v in packs.items() if not v["missing"] for p in (v["eix"], v["epk"])}
    stray = sorted(p.name for p in pack_dir.iterdir() if p.name.lower() not in known and p.name != "Index")
    return pack_dir, keys, order, packs, stray


# ---------------------------------------------------------------- commands

def decoded_sizes(p: dict) -> dict[int, int]:
    """data_position -> decoded bytes (CLZObject THeader.dwRealSize for types 1/2, data_size otherwise)."""
    out = {}
    with open(p["epk"], "rb") as f:
        for e in p["entries"]:
            if e["type"] in (1, 2):
                f.seek(e["pos"])
                fourcc, _, _, real = struct.unpack("<4sIII", f.read(16))
                if fourcc != MCOZ:
                    raise ValueError(f"{e['name']}: no MCOZ header")
                out[e["pos"]] = real
            else:
                out[e["pos"]] = e["size"]
    return out


def summary(args) -> dict:
    client = Path(args.client) if args.client else refroot.reference_root().parents[1] / "Client"
    pack_dir, keys, order, packs, stray = load(client, refroot.reference_root())

    per_type = defaultdict(lambda: {"files": 0, "stored_bytes": 0, "decoded_bytes": 0})
    per_pack = {}
    crc_mismatch = 0
    not_lower = 0
    backslash = 0
    prefixes = Counter()
    dict_first: dict[str, tuple[str, dict]] = {}
    overridden = []
    for name, _ in order:
        p = packs[name]
        if p["missing"] or name in per_pack:
            continue
        counts = Counter()
        size = Counter()
        decoded = decoded_sizes(p)
        for e in p["entries"]:
            t = TYPE_NAMES.get(e["type"], str(e["type"]))
            counts[t] += 1
            size[t] += e["size"]
            agg = per_type[t]
            agg["files"] += 1
            agg["stored_bytes"] += e["size"]
            agg["decoded_bytes"] += decoded[e["pos"]]
            fn = e["name"]
            if zlib.crc32(fn.encode("latin-1")) != e["crc"]:
                crc_mismatch += 1
            if any("A" <= c <= "Z" for c in fn):  # ConvertFileName is an ASCII tolower; cp949 bytes pass through
                not_lower += 1
            if "\\" in fn:
                backslash += 1
            prefixes[fn.split("/")[0] + "/" + (fn.split("/")[1] + "/" if fn.count("/") > 1 else "")] += 1
            # CEterFileDict::InsertItem into an unordered_multimap; GetItem returns the first equal element.
            if fn in dict_first:
                w_pack, w = dict_first[fn]
                same = (w["size"], w["data_crc"]) == (e["size"], e["data_crc"])
                overridden.append({"file": fn, "kept": w_pack, "shadowed": name, "same_content": same})
            else:
                dict_first[fn] = (name, e)
        per_pack[name] = {"folder": p["folder"], "files": sum(counts.values()), "epk_bytes": p["epk_bytes"],
                          "index_encrypted": p["index_encrypted"], "types": dict(counts), "type_bytes": dict(size)}

    lower_groups = defaultdict(set)
    for fn in dict_first:
        lower_groups[fn.lower()].add(fn)
    case_collisions = sorted(v for v in (sorted(g) for g in lower_groups.values()) if len(v) > 1)

    report = {
        "pack_dir": "Client/pack",  # relative to the 40250 tree; no machine paths in committed reports
        "registered": [n for n, _ in order],
        "missing": sorted(n for n, v in packs.items() if v["missing"]),
        "stray_files": stray,
        "per_type": dict(per_type),
        "per_pack": per_pack,
        "unique_files": len(dict_first),
        "overrides": {"count": len(overridden), "different_content": sum(1 for o in overridden if not o["same_content"]),
                      "by_pair": Counter(f'{o["kept"]} > {o["shadowed"]}' for o in overridden).most_common(),
                      "different_examples": [o for o in overridden if not o["same_content"]][:20]},
        "paths": {"filename_crc_mismatch": crc_mismatch, "not_lowercase": not_lower, "backslash": backslash,
                  "case_collisions": case_collisions, "top_prefixes": prefixes.most_common(25)},
    }
    print_summary(report)
    if args.json:
        Path(args.json).write_text(json.dumps(report, indent=1, ensure_ascii=False) + "\n", encoding="utf-8")
        print(f"\nwrote {args.json}")
    return report


def print_summary(r: dict) -> None:
    print(f"pack dir: {r['pack_dir']} (40250 tree)")
    print(f"registered: {len(r['registered'])} (Index pairs x2 with _texcache, + root); "
          f"missing on disk: {len(r['missing'])}; stray files: {r['stray_files']}")
    print("\nper compressed_type:")
    for t, v in sorted(r["per_type"].items()):
        print(f"  {t:<22} {v['files']:>7} files  {v['stored_bytes'] / 2**20:>9.1f} MiB stored"
              f"  {v['decoded_bytes'] / 2**20:>9.1f} MiB decoded")
    print(f"\nunique paths: {r['unique_files']}; overridden entries: {r['overrides']['count']}"
          f" ({r['overrides']['different_content']} with different content)")
    for pair, n in r["overrides"]["by_pair"][:15]:
        print(f"  {n:>6}  {pair}")
    p = r["paths"]
    print(f"\npaths: crc mismatch {p['filename_crc_mismatch']}, not lowercase {p['not_lowercase']}, "
          f"backslash {p['backslash']}, case collisions {len(p['case_collisions'])}")
    for pre, n in p["top_prefixes"][:12]:
        print(f"  {n:>7}  {pre}")
    print("\nper pack (types):")
    for name, v in r["per_pack"].items():
        ts = ", ".join(f"{t}={n}" for t, n in sorted(v["types"].items()))
        print(f"  {name:<40} {v['folder']:<14} {v['files']:>6}  {v['epk_bytes'] / 2**20:>7.1f} MiB  {ts}")


def lz_object_len(head: bytes) -> int:
    """Bytes a CLZObject really occupies: THeader + (TEA block when encrypted, else fourcc + LZO)."""
    fourcc, enc, comp, _ = struct.unpack_from("<4sIII", head)
    if fourcc != MCOZ:
        raise ValueError("no MCOZ header")
    return 16 + (enc if enc else 4 + comp)


def verify(args) -> int:
    """Decode every sampled NONE/COMPRESS/SECURITY entry and classify how its pack was written.

    Two packers wrote these files. The Ymir one stores data_size >= the whole CLZObject (often padded to a
    256-byte multiple) and data_crc over the stored bytes, which is what CEterPack::Get2 checks for SECURITY/PANAMA. The other ("short") stores
    data_size 16 bytes short of the object and data_crc over the decoded bytes: 40250 still decodes its
    COMPRESS entries (CLZO::Decompress trusts the header and reads past data_size inside the mapped view) but
    rejects its SECURITY entries in the CRC check, so those files never load in the original client.
    """
    client = Path(args.client) if args.client else refroot.reference_root().parents[1] / "Client"
    _, keys, _, packs, _ = load(client, refroot.reference_root())
    ok = bad = 0
    layout = Counter()
    crc_stats = Counter()
    rejected = []
    for name, p in packs.items():
        if p["missing"]:
            continue
        by_type = defaultdict(list)
        for e in p["entries"]:
            by_type[e["type"]].append(e)
        with open(p["epk"], "rb") as f:
            for t, es in sorted(by_type.items()):
                if t not in LOCAL_TYPES:
                    continue
                es = sorted(es, key=lambda e: e["size"])
                picks = {id(e): e for e in (es[:: max(1, len(es) // args.per_type)][: args.per_type] + [es[-1]])}
                for e in picks.values():
                    try:
                        f.seek(e["pos"])
                        stored = f.read(e["size"])
                        if len(stored) != e["size"]:
                            raise ValueError("short read")
                        if t == 0:
                            out, kind = stored, "none"
                        else:
                            f.seek(e["pos"])
                            obj = f.read(lz_object_len(stored[:16]))
                            gap = len(obj) - e["size"]
                            kind = "ymir" if gap <= 0 else "short" if gap == 16 else None
                            if kind is None:
                                raise ValueError(f"data_size {e['size']} vs object {len(obj)}")
                            out = lz_object(obj, keys["s_adwEterPackSecurityKey"] if t == 2 else None)
                        # Only SECURITY/PANAMA are CRC-checked by the client; tally what the CRC covers.
                        crc = e["data_crc"]
                        covers = ("stored" if zlib.crc32(stored) == crc else
                                  "decoded" if zlib.crc32(out) == crc else "neither")
                        layout[f"{name}:{kind}"] += 1
                        crc_stats[(TYPE_NAMES[t], kind, covers)] += 1
                        if t == 2 and zlib.crc32(stored) != e["data_crc"]:
                            rejected.append(f"{name}:{e['name']}")
                        ok += 1
                    except (ValueError, IndexError, struct.error) as exc:
                        bad += 1
                        print(f"FAIL {name}:{e['name']} ({TYPE_NAMES[t]}): {exc}")
    short = sorted({k.split(":")[0] for k in layout if k.endswith(":short")})
    print(f"short-layout packs: {short}")
    print("data_crc covers (type, layout, covers): " + ", ".join(f"{k}={v}" for k, v in sorted(crc_stats.items())))
    print(f"SECURITY entries the 40250 CRC check rejects: {len(rejected)}")
    for r in rejected:
        print(f"  {r}")
    print(f"verify: {ok} decoded, {bad} failed")
    return 1 if bad else 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--client", default=os.environ.get("MT_40250_CLIENT"),
                    help="40250 Client dir (default: <reference root>/../../Client)")
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("summary")
    s.add_argument("--json")
    v = sub.add_parser("verify")
    v.add_argument("--per-type", type=int, default=3)
    args = ap.parse_args()
    if args.cmd == "summary":
        summary(args)
        return 0
    return verify(args)


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