#!/usr/bin/env python3
# BEGIN CHANGE: patch items.otb clientIds (server id fixed)
"""
Patch client IDs in items.otb by server item_id.

Usage:
  python otb_patch_client_ids.py --otb PATH --map '{"2160":3043}' --backup
  python otb_patch_client_ids.py --otb PATH --set 2160=9999 --dry-run
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import struct
import sys
import time
from pathlib import Path

ESCAPE, NODE_START, NODE_END = 0xFD, 0xFE, 0xFF
ATTR_SERVER, ATTR_CLIENT = 16, 17


def read_node_buffer(raw: bytes, start: int) -> tuple[bytes, int]:
    buf = bytearray()
    p = start
    n = len(raw)
    while p < n:
        b = raw[p]
        p += 1
        if b == NODE_START:
            depth = 1
            while p < n and depth > 0:
                nb = raw[p]
                p += 1
                if nb == NODE_START:
                    depth += 1
                elif nb == NODE_END:
                    depth -= 1
                elif nb == ESCAPE and p < n:
                    p += 1
        elif b == NODE_END:
            break
        elif b == ESCAPE:
            if p >= n:
                break
            buf.append(raw[p])
            p += 1
        else:
            buf.append(b)
    return bytes(buf), p


def escape_bytes(data: bytes) -> bytes:
    out = bytearray()
    for b in data:
        if b in (ESCAPE, NODE_START, NODE_END):
            out.append(ESCAPE)
        out.append(b)
    return bytes(out)


def parse_otb(raw: bytes) -> tuple[bytes, list[bytearray]]:
    if len(raw) < 8 or raw[0:4] != b"\x00\x00\x00\x00":
        raise ValueError("invalid otb signature")
    if raw[4] != NODE_START:
        raise ValueError("missing root NODE_START")
    root_start = 5
    # root buffer (attrs only) + children collected like getChildren
    # First pass: walk from root_start for children markers interleaved with root payload
    # Same algorithm as portal items-otb-lib / extract probe
    children: list[bytearray] = []
    p = root_start
    root_buf = bytearray()
    n = len(raw)
    # Reconstruct: unserialize root node into root_buf while collecting child starts
    # Mirror BinaryTree.unserialize + getChildren:
    # For rewrite we only need: rootAttrs buffer and list of child buffers.
    # getChildren walks from startPos of root (after NODE_START of root):
    p = root_start
    while p < n:
        b = raw[p]
        p += 1
        if b == NODE_START:
            cbuf, p = read_node_buffer(raw, p)
            children.append(bytearray(cbuf))
        elif b == NODE_END:
            break
        elif b == ESCAPE:
            if p < n:
                # root attr payload byte (escaped)
                # Actually for root, attrs appear BEFORE children in file stream
                # mixed: root data bytes and child nodes. unserialize puts non-node bytes in buffer.
                root_buf.append(raw[p])
                p += 1
        else:
            root_buf.append(b)

    # Problem: the walk above treats ALL non-node as root_buf, including wrongly if structure differs.
    # OTL structure: root node content = attrs; children are nested. unserialize skips nested
    # by skipNodes without adding. getChildren finds NODE_STARTs at root level.
    # Our walk matches getChildren for children; for root_buf we need unserialize path.

    # Re-parse root_buf properly with unserialize semantics:
    root_buf2 = bytearray()
    p = root_start
    while p < n:
        b = raw[p]
        p += 1
        if b == NODE_START:
            # skip nested child entirely
            depth = 1
            while p < n and depth > 0:
                nb = raw[p]
                p += 1
                if nb == NODE_START:
                    depth += 1
                elif nb == NODE_END:
                    depth -= 1
                elif nb == ESCAPE and p < n:
                    p += 1
        elif b == NODE_END:
            break
        elif b == ESCAPE:
            if p < n:
                root_buf2.append(raw[p])
                p += 1
        else:
            root_buf2.append(b)

    children = []
    p = root_start
    while p < n:
        b = raw[p]
        p += 1
        if b == NODE_START:
            cbuf, p = read_node_buffer(raw, p)
            children.append(bytearray(cbuf))
        elif b == NODE_END:
            break
        elif b == ESCAPE:
            if p < n:
                p += 1
        # else: root payload, skip

    return bytes(root_buf2), children


def patch_child_client_id(buf: bytearray, new_cid: int) -> bool:
    """Return True if server id matched path is handled by caller; patch client attr in buf."""
    if len(buf) < 5:
        return False
    i = 5  # skip category + flags
    found_client = False
    while i < len(buf):
        attr = buf[i]
        i += 1
        if i + 2 > len(buf):
            break
        alen = buf[i] | (buf[i + 1] << 8)
        i += 2
        if i + alen > len(buf):
            break
        if attr == ATTR_CLIENT and alen >= 2:
            struct.pack_into("<H", buf, i, int(new_cid) & 0xFFFF)
            found_client = True
        i += alen
    if not found_client:
        # append client id attr
        buf.extend(struct.pack("<BHH", ATTR_CLIENT, 2, int(new_cid) & 0xFFFF))
    return True


def child_server_id(buf: bytes) -> int | None:
    if len(buf) < 5:
        return None
    i = 5
    while i < len(buf):
        attr = buf[i]
        i += 1
        if i + 2 > len(buf):
            break
        alen = buf[i] | (buf[i + 1] << 8)
        i += 2
        if i + alen > len(buf):
            break
        if attr == ATTR_SERVER and alen >= 2:
            return buf[i] | (buf[i + 1] << 8)
        i += alen
    return None


def build_otb(root_buf: bytes, children: list[bytearray]) -> bytes:
    out = bytearray()
    out += b"\x00\x00\x00\x00"
    out.append(NODE_START)
    out += escape_bytes(root_buf)
    for ch in children:
        out.append(NODE_START)
        out += escape_bytes(bytes(ch))
        out.append(NODE_END)
    out.append(NODE_END)
    return bytes(out)


def apply_patches(otb_path: Path, mapping: dict[int, int], backup: bool, dry_run: bool) -> dict:
    raw = otb_path.read_bytes()
    root_buf, children = parse_otb(raw)
    patched = 0
    missing = []
    for sid, cid in mapping.items():
        sid, cid = int(sid), int(cid)
        hit = False
        for ch in children:
            s = child_server_id(bytes(ch))
            if s == sid:
                patch_child_client_id(ch, cid)
                patched += 1
                hit = True
                break
        if not hit:
            missing.append(sid)

    out = build_otb(root_buf, children)
    # verify parse of result
    from_map = {}
    # quick re-parse server->client
    _, ch2 = parse_otb(out)
    for ch in ch2:
        sid = child_server_id(bytes(ch))
        if sid is None:
            continue
        # read client
        i = 5
        cid = None
        buf = bytes(ch)
        while i < len(buf):
            attr = buf[i]
            i += 1
            if i + 2 > len(buf):
                break
            alen = buf[i] | (buf[i + 1] << 8)
            i += 2
            if i + alen > len(buf):
                break
            if attr == ATTR_CLIENT and alen >= 2:
                cid = buf[i] | (buf[i + 1] << 8)
            i += alen
        if cid is not None:
            from_map[sid] = cid

    for sid, cid in mapping.items():
        if int(sid) in from_map and from_map[int(sid)] != int(cid):
            raise RuntimeError("verify failed for server %s" % sid)

    result = {
        "ok": True,
        "patched": patched,
        "missing": missing,
        "bytes_in": len(raw),
        "bytes_out": len(out),
        "children": len(children),
    }
    if dry_run:
        result["dry_run"] = True
        return result

    if backup:
        bak = otb_path.with_name(
            otb_path.name + ".bak_" + time.strftime("%Y%m%d_%H%M%S")
        )
        shutil.copy2(otb_path, bak)
        result["backup"] = str(bak)

    # Atomic replace: works even when the existing file is root:root 644,
    # as long as the directory is writable (www-data can unlink+rename).
    tmp = otb_path.with_name(otb_path.name + ".tmp_write")
    tmp.write_bytes(out)
    os.replace(tmp, otb_path)
    try:
        os.chmod(otb_path, 0o664)
    except OSError:
        pass
    result["otb"] = str(otb_path)
    return result


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--otb", type=Path, required=True)
    ap.add_argument("--map", type=str, default="", help='JSON object {"2160":1234}')
    ap.add_argument("--set", action="append", default=[], help="server=client (repeatable)")
    ap.add_argument("--backup", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()
    mapping: dict[int, int] = {}
    if args.map.strip():
        obj = json.loads(args.map)
        for k, v in obj.items():
            mapping[int(k)] = int(v)
    for s in args.set:
        a, b = s.split("=", 1)
        mapping[int(a)] = int(b)
    if not mapping:
        print(json.dumps({"ok": False, "error": "empty mapping"}))
        return 1
    if not args.otb.is_file():
        print(json.dumps({"ok": False, "error": "otb missing"}))
        return 1
    try:
        result = apply_patches(args.otb, mapping, backup=args.backup, dry_run=args.dry_run)
    except Exception as e:
        print(json.dumps({"ok": False, "error": str(e)}))
        return 1
    print(json.dumps(result))
    return 0 if result.get("ok") else 1


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