#!/usr/bin/env python3
"""
Phantasy Star Nova canonical RMD source builder.

Builds a deterministic scripts/ directory for psnova-textinserter from:
  1. the base-game NOVA_FileList_Vita.cpk (APP)
  2. the runtime/update NOVA_FileList_Vita_patch.cpk (PATCH)
  3. the PSNovaTranslations repository JSON records

Resolution rule:
    PATCH msg_<ID>.rmd if present, otherwise APP msg_<ID>.rmd.

Only repository IDs are emitted. CRILAYLA-packed CPK entries are decompressed
back to their RMD form before writing to scripts/.

This tool does not modify the game image or build a distributable patch.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import struct
import sys
import tempfile
import zipfile
from dataclasses import dataclass
from typing import BinaryIO, Dict, Iterable, List, Optional, Sequence, Tuple

RMD_RE = re.compile(r"^msg_(.+)\.rmd$", re.IGNORECASE)


class CPKError(RuntimeError):
    pass


def sha256_file(path: Path, block_size: int = 8 * 1024 * 1024) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            block = f.read(block_size)
            if not block:
                break
            h.update(block)
    return h.hexdigest()


def cri_utf_decrypt(data: bytes) -> bytes:
    """Decrypt a CRI encrypted @UTF table."""
    out = bytearray(data)
    key = 0x655F
    for i in range(len(out)):
        out[i] ^= key & 0xFF
        key = (key * 0x4115) & 0xFFFFFFFF
    return bytes(out)


def _c_string(buf: bytes, offset: int) -> str:
    if offset < 0 or offset >= len(buf):
        raise CPKError(f"String offset out of range: {offset}")
    end = buf.find(b"\x00", offset)
    if end < 0:
        end = len(buf)
    return buf[offset:end].decode("utf-8", "replace")


def _parse_utf_blob(raw: bytes) -> Tuple[List[dict], List[str]]:
    if raw[:4] != b"@UTF":
        raise CPKError("Decoded table does not begin with @UTF")
    if len(raw) < 0x20:
        raise CPKError("Truncated @UTF table")

    table_size = struct.unpack_from(">I", raw, 4)[0]
    body = raw[8:8 + table_size]
    if len(body) < 0x18:
        raise CPKError("Truncated @UTF body")

    rows_off, strings_off, data_off, _table_name_off, col_count, row_width, row_count = struct.unpack_from(
        ">IIIIHHI", body, 0
    )

    pos = 0x18
    columns: List[Tuple[str, int, int, object]] = []

    def read_value(p: int, typ: int) -> Tuple[object, int]:
        scalar = {
            0x0: ">B",
            0x1: ">b",
            0x2: ">H",
            0x3: ">h",
            0x4: ">I",
            0x5: ">i",
            0x6: ">Q",
            0x7: ">q",
            0x8: ">f",
            0x9: ">d",
        }
        if typ in scalar:
            fmt = scalar[typ]
            size = struct.calcsize(fmt)
            if p + size > len(body):
                raise CPKError("Value runs beyond @UTF body")
            return struct.unpack_from(fmt, body, p)[0], p + size
        if typ == 0xA:  # string
            off = struct.unpack_from(">I", body, p)[0]
            return _c_string(body, strings_off + off), p + 4
        if typ == 0xB:  # byte array
            off, size = struct.unpack_from(">II", body, p)
            start = data_off + off
            return body[start:start + size], p + 8
        raise CPKError(f"Unsupported @UTF value type 0x{typ:X}")

    for _ in range(col_count):
        flag = body[pos]
        pos += 1
        name_off = struct.unpack_from(">I", body, pos)[0]
        pos += 4
        name = _c_string(body, strings_off + name_off)
        storage = flag & 0xF0
        typ = flag & 0x0F
        const = None
        if storage == 0x30:
            const, pos = read_value(pos, typ)
        columns.append((name, storage, typ, const))

    rows: List[dict] = []
    for row_index in range(row_count):
        p = rows_off + row_index * row_width
        record = {}
        for name, storage, typ, const in columns:
            if storage == 0x50:
                value, p = read_value(p, typ)
            elif storage == 0x30:
                value = const
            elif storage == 0x10:
                value = None
            else:
                value = None
            record[name] = value
        rows.append(record)

    return rows, [c[0] for c in columns]


def read_utf_table(f: BinaryIO, offset: int) -> Tuple[List[dict], List[str]]:
    f.seek(offset)
    first8 = f.read(8)
    if len(first8) != 8:
        raise CPKError(f"Cannot read @UTF header at 0x{offset:X}")

    if first8[:4] == b"@UTF":
        table_size = struct.unpack_from(">I", first8, 4)[0]
        f.seek(offset)
        raw = f.read(8 + table_size)
    else:
        dec8 = cri_utf_decrypt(first8)
        if dec8[:4] != b"@UTF":
            raise CPKError(f"No plain/encrypted @UTF table at 0x{offset:X}")
        table_size = struct.unpack_from(">I", dec8, 4)[0]
        f.seek(offset)
        raw = cri_utf_decrypt(f.read(8 + table_size))

    return _parse_utf_blob(raw)


def crilayla_decompress(data: bytes) -> bytes:
    """Decompress a CRILAYLA payload to its original bytes."""
    if not data.startswith(b"CRILAYLA"):
        return data
    if len(data) < 0x10:
        raise CPKError("Truncated CRILAYLA header")

    uncompressed_size = struct.unpack_from("<I", data, 8)[0]
    header_offset = struct.unpack_from("<I", data, 12)[0]
    compressed_end = 0x10 + header_offset
    if compressed_end + 0x100 > len(data):
        raise CPKError("CRILAYLA header/data offsets are invalid")

    out = bytearray(uncompressed_size + 0x100)
    out[:0x100] = data[compressed_end:compressed_end + 0x100]

    in_pos = compressed_end - 1
    out_pos = len(out) - 1
    bit_pool = 0
    bits_left = 0

    def get_bits(count: int) -> int:
        nonlocal in_pos, bit_pool, bits_left
        result = 0
        while count:
            if bits_left == 0:
                if in_pos < 0:
                    raise CPKError("CRILAYLA bitstream ended unexpectedly")
                bit_pool = data[in_pos]
                in_pos -= 1
                bits_left = 8
            take = min(count, bits_left)
            shift = bits_left - take
            value = (bit_pool >> shift) & ((1 << take) - 1)
            result = (result << take) | value
            bits_left -= take
            bit_pool &= (1 << bits_left) - 1 if bits_left else 0
            count -= take
        return result

    while out_pos >= 0x100:
        if get_bits(1) == 0:
            out[out_pos] = get_bits(8)
            out_pos -= 1
            continue

        distance = get_bits(13) + 3
        length = 3
        for width in (2, 3, 5, 8):
            value = get_bits(width)
            length += value
            if value != (1 << width) - 1:
                break
        else:
            while True:
                value = get_bits(8)
                length += value
                if value != 0xFF:
                    break

        src_pos = out_pos + distance
        for _ in range(length):
            if out_pos < 0x100:
                break
            if src_pos >= len(out):
                raise CPKError("Invalid CRILAYLA back-reference")
            out[out_pos] = out[src_pos]
            out_pos -= 1
            src_pos -= 1

    return bytes(out)


@dataclass(frozen=True)
class CPKEntry:
    filename: str
    dirname: str
    file_size: int
    extract_size: int
    file_offset: int
    absolute_offset: int
    cpk_id: Optional[int]


class CPKIndex:
    def __init__(self, path: Path):
        self.path = Path(path)
        self.header: dict = {}
        self.entries: List[CPKEntry] = []
        self.rmd_by_id: Dict[str, CPKEntry] = {}
        self._load()

    def _load(self) -> None:
        with self.path.open("rb") as f:
            if f.read(4) != b"CPK ":
                raise CPKError(f"Not a CPK file: {self.path}")
            header_rows, _ = read_utf_table(f, 0x10)
            if not header_rows:
                raise CPKError("Empty CPK header table")
            self.header = header_rows[0]

            toc_offset = self.header.get("TocOffset")
            content_offset = self.header.get("ContentOffset")
            if toc_offset is None or content_offset is None:
                raise CPKError("CPK has no TOC/content offset")

            f.seek(int(toc_offset))
            if f.read(4) != b"TOC ":
                raise CPKError(f"Expected TOC chunk at 0x{int(toc_offset):X}")
            toc_rows, _ = read_utf_table(f, int(toc_offset) + 0x10)

            # CRI CPK TOC FileOffset values are relative to the lower of
            # ContentOffset and TocOffset for these Nova archives.
            base = min(int(content_offset), int(toc_offset))
            for row in toc_rows:
                filename = row.get("FileName") or ""
                dirname = row.get("DirName") or ""
                file_size = int(row.get("FileSize") or 0)
                extract_size = int(row.get("ExtractSize") or file_size)
                file_offset = int(row.get("FileOffset") or 0)
                entry = CPKEntry(
                    filename=filename,
                    dirname=dirname,
                    file_size=file_size,
                    extract_size=extract_size,
                    file_offset=file_offset,
                    absolute_offset=base + file_offset,
                    cpk_id=row.get("ID"),
                )
                self.entries.append(entry)
                m = RMD_RE.match(filename)
                if m:
                    rid = m.group(1)
                    old = self.rmd_by_id.get(rid)
                    if old is not None and old != entry:
                        raise CPKError(f"Duplicate RMD ID {rid!r} in {self.path}")
                    self.rmd_by_id[rid] = entry

    def read_entry(self, entry: CPKEntry, decompress: bool = True) -> bytes:
        with self.path.open("rb") as f:
            f.seek(entry.absolute_offset)
            raw = f.read(entry.file_size)
        if len(raw) != entry.file_size:
            raise CPKError(f"Short read for {entry.filename}")
        data = crilayla_decompress(raw) if decompress else raw
        if decompress and entry.extract_size and len(data) != entry.extract_size:
            raise CPKError(
                f"Extracted-size mismatch for {entry.filename}: "
                f"got {len(data)}, expected {entry.extract_size}"
            )
        return data


def repository_ids(repo_zip: Path) -> Tuple[List[str], Dict[str, str]]:
    ids: List[str] = []
    owners: Dict[str, str] = {}
    with zipfile.ZipFile(repo_zip, "r") as z:
        json_names = [
            n for n in z.namelist()
            if "/rmd/" in n.replace("\\", "/") and n.lower().endswith(".json")
        ]
        if not json_names:
            raise RuntimeError("No /rmd/*.json files found in repository ZIP")
        for name in sorted(json_names):
            data = json.loads(z.read(name).decode("utf-8-sig"))
            if not isinstance(data, dict):
                raise RuntimeError(f"Expected object in {name}")
            owner = Path(name).name
            for rid in data.keys():
                rid = str(rid)
                if rid in owners:
                    raise RuntimeError(
                        f"Duplicate repository ID {rid}: {owners[rid]} and {owner}"
                    )
                owners[rid] = owner
                ids.append(rid)
    return ids, owners


def find_patch_cpk_in_zip(patch_zip: Path, requested_member: Optional[str]) -> str:
    with zipfile.ZipFile(patch_zip, "r") as z:
        names = z.namelist()
        if requested_member:
            if requested_member not in names:
                raise RuntimeError(f"Patch ZIP member not found: {requested_member}")
            return requested_member
        candidates = [n for n in names if n.lower().endswith(".cpk") and "patch" in Path(n).name.lower()]
        if len(candidates) == 1:
            return candidates[0]
        exact = [n for n in names if Path(n).name.lower() == "nova_filelist_vita_patch.cpk"]
        if len(exact) == 1:
            return exact[0]
        raise RuntimeError(
            "Could not uniquely identify patch CPK inside ZIP. "
            "Use --patch-zip-member. Candidates: " + ", ".join(candidates[:20])
        )


def materialize_patch_cpk(
    patch_cpk: Optional[Path],
    patch_zip: Optional[Path],
    patch_zip_member: Optional[str],
    temp_dir: Path,
) -> Tuple[Path, str]:
    if patch_cpk:
        return Path(patch_cpk), str(Path(patch_cpk))
    if not patch_zip:
        raise RuntimeError("Supply either --patch-cpk or --patch-zip")
    member = find_patch_cpk_in_zip(Path(patch_zip), patch_zip_member)
    out = temp_dir / Path(member).name
    with zipfile.ZipFile(patch_zip, "r") as z, z.open(member) as src, out.open("wb") as dst:
        shutil.copyfileobj(src, dst, length=8 * 1024 * 1024)
    return out, f"{patch_zip}!/{member}"


def write_csv(path: Path, rows: Sequence[dict], fieldnames: Sequence[str]) -> None:
    with path.open("w", newline="", encoding="utf-8-sig") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        w.writerows(rows)


def build(args: argparse.Namespace) -> int:
    app_cpk_path = Path(args.app_cpk).resolve()
    repo_zip_path = Path(args.repo).resolve()
    out_root = Path(args.out).resolve()
    scripts_dir = out_root / "scripts"
    manifest_dir = out_root / "manifest"

    if not app_cpk_path.is_file():
        raise RuntimeError(f"APP CPK not found: {app_cpk_path}")
    if not repo_zip_path.is_file():
        raise RuntimeError(f"Repository ZIP not found: {repo_zip_path}")

    repo_ids, repo_owner = repository_ids(repo_zip_path)
    repo_set = set(repo_ids)

    with tempfile.TemporaryDirectory(prefix="psnova_canonical_") as td:
        temp_dir = Path(td)
        patch_path, patch_label = materialize_patch_cpk(
            Path(args.patch_cpk).resolve() if args.patch_cpk else None,
            Path(args.patch_zip).resolve() if args.patch_zip else None,
            args.patch_zip_member,
            temp_dir,
        )

        print(f"[1/5] Indexing APP:   {app_cpk_path}")
        app = CPKIndex(app_cpk_path)
        print(f"      files={len(app.entries):,} rmd_ids={len(app.rmd_by_id):,}")
        print(f"[2/5] Indexing PATCH: {patch_label}")
        patch = CPKIndex(patch_path)
        print(f"      files={len(patch.entries):,} rmd_ids={len(patch.rmd_by_id):,}")
        print(f"[3/5] Repository IDs: {len(repo_ids):,}")

        plan: List[dict] = []
        missing: List[str] = []
        counts = {"PATCH": 0, "APP": 0}
        projected_bytes = 0

        for rid in repo_ids:
            if rid in patch.rmd_by_id:
                source = "PATCH"
                entry = patch.rmd_by_id[rid]
            elif rid in app.rmd_by_id:
                source = "APP"
                entry = app.rmd_by_id[rid]
            else:
                missing.append(rid)
                continue
            counts[source] += 1
            projected_bytes += entry.extract_size
            plan.append({
                "record_id": rid,
                "repo_json": repo_owner[rid],
                "source": source,
                "source_cpk_file": entry.filename,
                "cpk_dir": entry.dirname,
                "cpk_id": "" if entry.cpk_id is None else entry.cpk_id,
                "file_offset": entry.file_offset,
                "absolute_offset": entry.absolute_offset,
                "packed_size": entry.file_size,
                "extract_size": entry.extract_size,
                "output_file": f"scripts/msg_{rid}.rmd",
            })

        if missing:
            sample = ", ".join(missing[:20])
            raise RuntimeError(
                f"Coverage failure: {len(missing)} repository IDs have no RMD in PATCH or APP. "
                f"First IDs: {sample}"
            )

        patch_extra = sorted(set(patch.rmd_by_id) - repo_set)
        app_extra = sorted(set(app.rmd_by_id) - repo_set)

        print(
            f"      resolution: PATCH={counts['PATCH']:,} APP={counts['APP']:,} missing=0"
        )
        print(f"      projected unpacked scripts size: {projected_bytes / (1024**2):,.1f} MiB")

        if args.dry_run:
            print("[4/5] Dry run requested. No RMD files written.")
            out_root.mkdir(parents=True, exist_ok=True)
            manifest_dir.mkdir(parents=True, exist_ok=True)
        else:
            if out_root.exists() and any(out_root.iterdir()) and not args.force and not args.resume:
                raise RuntimeError(
                    f"Output directory is not empty: {out_root}. Use --force to replace or --resume to continue."
                )
            if args.force:
                shutil.rmtree(scripts_dir, ignore_errors=True)
                shutil.rmtree(manifest_dir, ignore_errors=True)
            scripts_dir.mkdir(parents=True, exist_ok=True)
            manifest_dir.mkdir(parents=True, exist_ok=True)

            print(f"[4/5] Extracting {len(plan):,} canonical RMD templates...")
            app_f = app_cpk_path.open("rb")
            patch_f = patch_path.open("rb")
            try:
                for i, row in enumerate(plan, 1):
                    source_index = patch if row["source"] == "PATCH" else app
                    source_f = patch_f if row["source"] == "PATCH" else app_f
                    entry = source_index.rmd_by_id[row["record_id"]]
                    target = scripts_dir / f"msg_{row['record_id']}.rmd"
                    if args.resume and target.is_file() and target.stat().st_size == entry.extract_size:
                        if i % 1000 == 0 or i == len(plan):
                            print(f"      {i:,}/{len(plan):,} (resume)")
                        continue
                    source_f.seek(entry.absolute_offset)
                    packed = source_f.read(entry.file_size)
                    if len(packed) != entry.file_size:
                        raise CPKError(f"Short read for {entry.filename}")
                    unpacked = crilayla_decompress(packed)
                    if entry.extract_size and len(unpacked) != entry.extract_size:
                        raise CPKError(
                            f"Extracted-size mismatch {entry.filename}: {len(unpacked)} != {entry.extract_size}"
                        )
                    if len(unpacked) < 4 or unpacked[:4] not in (b" DMR", b"RMD ", b"RMD\x00"):
                        # Nova RMD observed magic is byte-reversed ' DMR'. Keep this as a
                        # warning rather than a hard error for unusual resource variants.
                        row["rmd_magic_warning"] = unpacked[:8].hex()
                    else:
                        row["rmd_magic_warning"] = ""
                    target.write_bytes(unpacked)
                    if args.hash_outputs:
                        row["output_sha256"] = hashlib.sha256(unpacked).hexdigest()
                    if i % 1000 == 0 or i == len(plan):
                        print(f"      {i:,}/{len(plan):,}")
            finally:
                app_f.close()
                patch_f.close()

        # Manifests are emitted in dry-run and full-build modes.
        for row in plan:
            row.setdefault("rmd_magic_warning", "")
            row.setdefault("output_sha256", "")

        fields = [
            "record_id", "repo_json", "source", "source_cpk_file", "cpk_dir", "cpk_id",
            "file_offset", "absolute_offset", "packed_size", "extract_size", "output_file",
            "rmd_magic_warning", "output_sha256",
        ]
        write_csv(manifest_dir / "canonical_source_manifest.csv", plan, fields)
        (manifest_dir / "patch_rmd_ids_not_in_repository.txt").write_text(
            "\n".join(patch_extra) + ("\n" if patch_extra else ""), encoding="utf-8"
        )
        (manifest_dir / "app_rmd_ids_not_in_repository.txt").write_text(
            "\n".join(app_extra) + ("\n" if app_extra else ""), encoding="utf-8"
        )
        (manifest_dir / "repository_ids_no_rmd.txt").write_text("", encoding="utf-8")

        summary = {
            "builder_version": "1.0.0",
            "resolution_rule": "PATCH_FIRST_APP_SECOND",
            "repository_records": len(repo_ids),
            "app_files": len(app.entries),
            "app_rmd_ids": len(app.rmd_by_id),
            "patch_files": len(patch.entries),
            "patch_rmd_ids": len(patch.rmd_by_id),
            "canonical_from_patch": counts["PATCH"],
            "canonical_from_app": counts["APP"],
            "repository_ids_no_rmd": 0,
            "projected_unpacked_bytes": projected_bytes,
            "patch_rmd_ids_not_in_repository": len(patch_extra),
            "app_rmd_ids_not_in_repository": len(app_extra),
            "dry_run": bool(args.dry_run),
            "inputs": {
                "app_cpk": str(app_cpk_path),
                "patch": patch_label,
                "repository_zip": str(repo_zip_path),
            },
        }
        if args.hash_inputs:
            print("[5/5] Hashing input artifacts...")
            summary["input_sha256"] = {
                "app_cpk": sha256_file(app_cpk_path),
                "patch_cpk": sha256_file(patch_path),
                "repository_zip": sha256_file(repo_zip_path),
            }
        else:
            print("[5/5] Writing manifests...")

        (manifest_dir / "SUMMARY.json").write_text(
            json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        text_summary = "\n".join([
            "PS NOVA CANONICAL SOURCE BUILDER",
            "================================",
            f"Repository records: {len(repo_ids):,}",
            f"APP RMD IDs:        {len(app.rmd_by_id):,}",
            f"PATCH RMD IDs:      {len(patch.rmd_by_id):,}",
            f"Canonical PATCH:    {counts['PATCH']:,}",
            f"Canonical APP:      {counts['APP']:,}",
            "Missing source RMD:  0",
            f"Unpacked scripts:   {projected_bytes / (1024**2):,.1f} MiB projected",
            "Resolution rule:     PATCH_FIRST_APP_SECOND",
            f"Mode:                {'DRY RUN' if args.dry_run else 'FULL BUILD'}",
            "",
        ])
        (manifest_dir / "SUMMARY.txt").write_text(text_summary, encoding="utf-8")

    print("DONE")
    print(f"Output: {out_root}")
    return 0


def make_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Build canonical PS Nova RMD source templates using PATCH-first, APP-second resolution."
    )
    p.add_argument("--app-cpk", required=True, help="Path to NOVA_FileList_Vita.cpk")
    patch = p.add_mutually_exclusive_group(required=True)
    patch.add_argument("--patch-cpk", help="Path to NOVA_FileList_Vita_patch.cpk")
    patch.add_argument("--patch-zip", help="ZIP containing NOVA_FileList_Vita_patch.cpk")
    p.add_argument("--patch-zip-member", help="Explicit inner CPK path when --patch-zip is ambiguous")
    p.add_argument("--repo", required=True, help="PSNovaTranslations repository ZIP")
    p.add_argument("--out", required=True, help="Output root; scripts/ and manifest/ are created here")
    p.add_argument("--dry-run", action="store_true", help="Index/resolve/manifest only; do not extract RMDs")
    p.add_argument("--force", action="store_true", help="Replace generated scripts/ and manifest/ if output exists")
    p.add_argument("--resume", action="store_true", help="Resume a partial full build; valid existing-size RMDs are skipped")
    p.add_argument("--hash-inputs", action="store_true", help="SHA-256 the large input CPKs (slower)")
    p.add_argument("--hash-outputs", action="store_true", help="SHA-256 every emitted RMD")
    return p


def main() -> int:
    args = make_parser().parse_args()
    try:
        return build(args)
    except (RuntimeError, CPKError, zipfile.BadZipFile, json.JSONDecodeError, OSError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
