"""Pure-Python PSP/PSX/theme/DLC PSN ``.pkg`` decryption (no Wine / pkg2zip).

Ports the folder-mode extraction used by pkgi-psp's depackager and pkg2zip:
parse the ``\\x7FPKG`` header, derive the PSP/PS3 AES key, AES-CTR-decrypt the
file entries, and write the game's ``EBOOT.PBP`` (+ siblings) into a folder.

Supported content:

* PSP games / minis (``EBOOT.PBP``) — ``DATA.PSP`` stays EDAT-encrypted, so it
  runs on CFW with ``npdrm_free`` (ARK ships it), same as pkgi-psp's install.
* PS1 Classics (``EBOOT.PBP`` + ``DOCUMENT.DAT``).
* PSP themes (``.ptf``) — the packaged PTF is EDAT-wrapped and is unwrapped here.
* PSP DLC (``.edat``) — extracted as-is into the owning game's folder.

Theme unwrapping needs ``pycryptodome`` (for AES-CMAC + ECB); everything else
works with either pycryptodome or the built-in pure-Python AES fallback.
"""
from __future__ import annotations

from pathlib import Path

PKG_MAGIC = 0x7F504B47

# Static PSN keys (key type 1 = PSP).
PSP_KEY = bytes.fromhex("07f2c68290b50d2c33818d709b60e62b")
PS3_KEY = bytes.fromhex("2e7b71d7c9c9a14ea3221f188828b8f8")

# PSP KIRK/AMCTL keys used to unwrap theme PTF EDATs (pkgi-psp / libkirk).
_KIRK38 = bytes.fromhex("12468d7e1c42209bba5426835eb03303")
_KIRK39 = bytes.fromhex("c43bb6d653ee67493ea95fbc0ced6f8a")
_KIRK63 = bytes.fromhex("9c9b1372f8c640cf1c62f5d592ddb582")
_AMCTL3 = bytes.fromhex("e350ed1d910a1fd029bb1c3ef34077fb")
_AMCTL4 = bytes.fromhex("135fa47cab395ba476b8cca98f3a0445")
_AMCTL5 = bytes.fromhex("678d7fa32a9ca0d1508ad8385e4b017e")

try:
    from Crypto.Cipher import AES as _AES  # pycryptodome (fast)
    _HAVE_CRYPTO = True
except ImportError:  # noqa: BLE001
    _HAVE_CRYPTO = False


# --------------------------------------------------------------------------- #
# Pure-Python AES-128 (fallback when pycryptodome is unavailable)
# --------------------------------------------------------------------------- #
_SBOX = bytes.fromhex(
    "637c777bf26b6fc53001672bfed7ab76ca82c97dfa5947f0add4a2af9ca472c0"
    "b7fd9326363ff7cc34a5e5f171d8311504c723c31896059a071280e2eb27b275"
    "09832c1a1b6e5aa0523bd6b329e32f8453d100ed20fcb15b6acbbe394a4c58cf"
    "d0efaafb434d338545f9027f503c9fa851a3408f929d38f5bcb6da2110fff3d2"
    "cd0c13ec5f974417c4a77e3d645d197360814fdc222a908846eeb814de5e0bdb"
    "e0323a0a4906245cc2d3ac629195e479e7c8376d8dd54ea96c56f4ea657aae08"
    "ba78252e1ca6b4c6e8dd741f4bbd8b8a703eb5664803f60e613557b986c11d9e"
    "e1f8981169d98e949b1e87e9ce5528df8ca1890dbfe6426841992d0fb054bb16"
)
_RCON = (0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36)


def _key_schedule(key):
    w = list(key)
    for i in range(4, 44):
        t = list(w[(i - 1) * 4:(i - 1) * 4 + 4])
        if i % 4 == 0:
            t = t[1:] + t[:1]
            t = [_SBOX[b] for b in t]
            t[0] ^= _RCON[i // 4 - 1]
        w += [w[(i - 4) * 4 + j] ^ t[j] for j in range(4)]
    return [bytes(w[i * 16:(i + 1) * 16]) for i in range(11)]


def _shift_rows(s):
    return [s[0], s[5], s[10], s[15], s[4], s[9], s[14], s[3],
            s[8], s[13], s[2], s[7], s[12], s[1], s[6], s[11]]


def _xtime(a):
    a <<= 1
    return (a ^ 0x1B) & 0xFF if a & 0x100 else a & 0xFF


def _mix_columns(s):
    out = [0] * 16
    for c in range(4):
        a0, a1, a2, a3 = s[4 * c], s[4 * c + 1], s[4 * c + 2], s[4 * c + 3]
        out[4 * c] = _xtime(a0) ^ (_xtime(a1) ^ a1) ^ a2 ^ a3
        out[4 * c + 1] = a0 ^ _xtime(a1) ^ (_xtime(a2) ^ a2) ^ a3
        out[4 * c + 2] = a0 ^ a1 ^ _xtime(a2) ^ (_xtime(a3) ^ a3)
        out[4 * c + 3] = (_xtime(a0) ^ a0) ^ a1 ^ a2 ^ _xtime(a3)
    return out


def _aes_encrypt_block(rk, block):
    s = [block[i] ^ rk[0][i] for i in range(16)]
    for rnd in range(1, 10):
        s = [_SBOX[b] for b in s]
        s = _shift_rows(s)
        s = _mix_columns(s)
        s = [s[i] ^ rk[rnd][i] for i in range(16)]
    s = [_SBOX[b] for b in s]
    s = _shift_rows(s)
    s = [s[i] ^ rk[10][i] for i in range(16)]
    return bytes(s)


def _add128(counter, n):
    c = bytearray(counter)
    for i in range(15, -1, -1):
        n += c[i]
        c[i] = n & 0xFF
        n >>= 8
    return bytes(c)


def _ctr_decrypt_py(key, iv, block_index, data):
    """AES-CTR decrypt (pure-Python): counter = iv + block_index, +1 per block."""
    rk = _key_schedule(key)
    counter = _add128(iv, block_index)
    out = bytearray()
    for i in range(0, len(data), 16):
        ks = _aes_encrypt_block(rk, counter)
        out += bytes(b ^ k for b, k in zip(data[i:i + 16], ks))
        counter = _add128(counter, 1)
    return bytes(out)


def _ctr_decrypt(key, iv, block_index, data):
    """AES-CTR decrypt: 128-bit big-endian counter at ``iv + block_index``."""
    if _HAVE_CRYPTO:
        nonce = iv[:8]
        init = int.from_bytes(iv[8:16], "big") + block_index
        return _AES.new(key, _AES.MODE_CTR, nonce=nonce, initial_value=init).decrypt(data)
    return _ctr_decrypt_py(key, iv, block_index, data)


def _make_ctr_decryptor(key, iv, block_index):
    """Return a streaming decryptor (callable: bytes -> bytes) for a data run."""
    if _HAVE_CRYPTO:
        nonce = iv[:8]
        init = int.from_bytes(iv[8:16], "big") + block_index
        cipher = _AES.new(key, _AES.MODE_CTR, nonce=nonce, initial_value=init)
        return cipher.decrypt
    rk = _key_schedule(key)
    counter = _add128(iv, block_index)
    state = {"rk": rk, "counter": counter}

    def decrypt(chunk):
        out = bytearray()
        for i in range(0, len(chunk), 16):
            ks = _aes_encrypt_block(state["rk"], state["counter"])
            out += bytes(b ^ k for b, k in zip(chunk[i:i + 16], ks))
            state["counter"] = _add128(state["counter"], 1)
        return bytes(out)
    return decrypt


# --------------------------------------------------------------------------- #
# PSP theme PTF EDAT unwrapping (pkgi-psp's unpack_psp_edat)
# --------------------------------------------------------------------------- #
def _need_crypto():
    if not _HAVE_CRYPTO:
        raise RuntimeError("theme extraction needs pycryptodome (pip install pycryptodome)")


def _ecb_decrypt(key, block):
    _need_crypto()
    return _AES.new(key, _AES.MODE_ECB).decrypt(block)


def _cmac(key, data):
    _need_crypto()
    from Crypto.Hash import CMAC
    c = CMAC.new(key, ciphermod=_AES)
    c.update(data)
    return c.digest()


def _kirk_iv(mac, key_header, offset2):
    tmp = _ecb_decrypt(_KIRK38, bytes(key_header[0x70:0x80]))
    iv = bytes(mac[i] ^ tmp[i] ^ key_header[offset2 + i] ^ _AMCTL3[i] ^ _AMCTL5[i] for i in range(16))
    iv = _ecb_decrypt(_KIRK39, iv)
    return bytes(b ^ _AMCTL4[i] for i, b in enumerate(iv))


def _kirk_cbc_decrypt(iv, index, data):
    prev = bytearray(16) if index == 0 else bytearray(iv[:12]) + index.to_bytes(4, "little")
    block = bytearray(iv[:12]) + index.to_bytes(4, "little")
    out = bytearray()
    for i in range(0, len(data), 16):
        v = int.from_bytes(block[12:16], "little") + 1
        block[12:16] = v.to_bytes(4, "little")
        o = _ecb_decrypt(_KIRK63, bytes(block))
        out += bytes(b ^ prev[k] ^ o[k] for k, b in enumerate(data[i:i + 16]))
        prev = bytes(block)
    return bytes(out)


def _unwrap_theme_ptf(buf):
    """Unwrap a PSP theme PTF EDAT (``buf`` is the fully CTR-decrypted entry)."""
    if len(buf) < 0x90 + 0xA0:
        raise RuntimeError("theme EDAT is too short")
    kh_off = buf[0xC]
    key_header = bytearray(buf[kh_off:kh_off + 0xA0])
    if bytes(key_header[:4]) != b"\x00PGD":
        raise RuntimeError("theme EDAT: bad PGD header")
    if int.from_bytes(key_header[4:8], "little") != 1 or int.from_bytes(key_header[8:12], "little") != 1:
        raise RuntimeError("theme EDAT: unsupported key/drm type")

    mac = _cmac(_KIRK38, bytes(key_header[:0x70]))

    # Step 1: reveal the data size/offset fields (bytes 0x30..0x60).
    key_header[0x30:0x60] = _kirk_cbc_decrypt(_kirk_iv(mac, key_header, 0x10), 0, bytes(key_header[0x30:0x60]))
    data_size = int.from_bytes(key_header[0x44:0x48], "little")
    data_offset = int.from_bytes(key_header[0x4c:0x50], "little")
    if data_offset != 0x90:
        raise RuntimeError("theme EDAT: unexpected data offset")

    # Step 2: decrypt the PTF body.
    iv = _kirk_iv(mac, key_header, 0x30)
    out = bytearray()
    nblocks = (data_size + 15) // 16
    for i in range(nblocks):
        chunk = buf[kh_off + data_offset + i * 16:kh_off + data_offset + (i + 1) * 16]
        out += _kirk_cbc_decrypt(iv, i, chunk)
    return bytes(out[:data_size])


# --------------------------------------------------------------------------- #
# PKG extraction
# --------------------------------------------------------------------------- #
def extract(pkg_path, dest_dir, kind_hint=None, progress_cb=None, cancel_event=None):
    """Decrypt a PSN pkg into ``dest_dir/<TITLE_ID>/``.

    Returns ``(kind, title_id, folder)`` where kind is ``"psp"``, ``"psx"``,
    ``"theme"`` or ``"dlc"``. ``kind_hint`` ("theme"/"dlc") overrides the
    pkg-header content-type detection (NoPayStation DLC ships as two formats).
    """
    from .store import TransferCancelled

    pkg_path = Path(pkg_path)
    dest_dir = Path(dest_dir)

    with open(pkg_path, "rb") as fh:
        header = fh.read(0x100)
    if len(header) < 0xE8 or int.from_bytes(header[0:4], "big") != PKG_MAGIC:
        raise RuntimeError(f"not a PKG file: {pkg_path.name}")

    meta_offset = int.from_bytes(header[8:12], "big")
    meta_count = int.from_bytes(header[12:16], "big")
    item_count = int.from_bytes(header[20:24], "big")
    enc_offset = int.from_bytes(header[32:40], "big")
    iv = header[0x70:0x80]
    key_type = header[0xE7] & 7
    title_id = header[0x37:0x40].decode("ascii", "replace").rstrip("\x00").strip() or "PKG"

    content_type = 0
    items_offset = 0
    with open(pkg_path, "rb") as fh:
        off = meta_offset
        for _ in range(meta_count):
            fh.seek(off)
            m = fh.read(16)
            if len(m) < 16:
                break
            mtype = int.from_bytes(m[0:4], "big")
            msize = int.from_bytes(m[4:8], "big")
            if mtype == 2:
                content_type = int.from_bytes(m[8:12], "big")
            elif mtype == 13:
                items_offset = int.from_bytes(m[8:12], "big")
            off += 8 + msize

    if kind_hint in ("theme", "dlc"):
        kind = kind_hint
    elif content_type == 6:
        kind = "psx"
    elif content_type in (7, 0xE, 0xF, 0x10):
        kind = "psp"
    elif content_type == 9:
        kind = "theme"
    elif content_type == 0xB:
        kind = "dlc"
    else:
        raise RuntimeError(f"unsupported NoPayStation content type 0x{content_type:x}")

    # NoPayStation DLC is wrapped in a PS3-format pkg (key type 4 in the header
    # but actually encrypted with the PS3 key); everything else uses the PSP key.
    if content_type == 0xB:
        main_key = PS3_KEY
        ps3_key = PS3_KEY
    elif key_type == 1:
        main_key = PSP_KEY
        ps3_key = PS3_KEY
    else:
        raise RuntimeError(f"unsupported PKG key type {key_type}")

    # Decrypt the whole item table in one shot.
    with open(pkg_path, "rb") as fh:
        fh.seek(enc_offset + items_offset)
        table = fh.read(item_count * 32)
    table = _ctr_decrypt(main_key, iv, items_offset // 16, table)

    entries = []
    for i in range(item_count):
        e = table[i * 32:(i + 1) * 32]
        name_offset = int.from_bytes(e[0:4], "big")
        name_size = int.from_bytes(e[4:8], "big")
        data_offset = int.from_bytes(e[8:16], "big")
        data_size = int.from_bytes(e[16:24], "big")
        psp_type = e[24]
        flags = e[27]
        if flags in (4, 18) or data_size == 0:
            continue
        item_key = main_key if psp_type == 0x90 else ps3_key
        entries.append((name_offset, name_size, data_offset, data_size, item_key))

    folder = dest_dir / title_id
    folder.mkdir(parents=True, exist_ok=True)
    written = []

    total = sum(ds for *_, ds, _ in entries) or 1
    done = 0

    with open(pkg_path, "rb") as fh:
        for name_offset, name_size, data_offset, data_size, item_key in entries:
            if cancel_event is not None and cancel_event.is_set():
                raise TransferCancelled()
            fh.seek(enc_offset + name_offset)
            name = _ctr_decrypt(item_key, iv, name_offset // 16, fh.read(name_size))
            name = name.split(b"\x00", 1)[0].decode("utf-8", "replace")
            base = name.rsplit("/", 1)[-1]

            if kind == "psx" and base not in ("EBOOT.PBP", "DOCUMENT.DAT"):
                continue
            if kind == "psp" and base not in ("EBOOT.PBP", "PSP-KEY.EDAT", "CONTENT.DAT"):
                continue
            if kind == "theme" and not base.lower().endswith(".ptf"):
                continue
            if kind == "dlc" and not base.lower().endswith(".edat"):
                continue
            if kind == "dlc" and base.upper() == "ISO.BIN.EDAT":
                # Small EDAT wrapper shipped inside PSP DLC packages — not DLC
                # content itself, so don't install it (matches pkgi-psp).
                continue

            if kind == "theme":
                # Read the (small) entry fully, CTR-decrypt, then unwrap the EDAT.
                fh.seek(enc_offset + data_offset)
                raw = _ctr_decrypt(item_key, iv, data_offset // 16, fh.read(data_size))
                dest = folder / base
                dest.write_bytes(_unwrap_theme_ptf(raw))
                done += data_size
            else:
                dest = folder / base
                decrypt = _make_ctr_decryptor(item_key, iv, data_offset // 16)
                fh.seek(enc_offset + data_offset)
                with open(dest, "wb") as out:
                    remaining = data_size
                    while remaining > 0:
                        if cancel_event is not None and cancel_event.is_set():
                            raise TransferCancelled()
                        chunk = fh.read(min(1 << 20, remaining))
                        if not chunk:
                            break
                        out.write(decrypt(chunk))
                        remaining -= len(chunk)
                        done += len(chunk)
            if progress_cb is not None:
                progress_cb("Extracting", min(done / total, 1.0), "")
            written.append(dest)

    if kind in ("psp", "psx") and not (folder / "EBOOT.PBP").exists():
        raise RuntimeError(f"no EBOOT.PBP found in {pkg_path.name}")
    if kind == "theme" and not any(p.suffix.lower() == ".ptf" for p in written):
        raise RuntimeError(f"no PTF found in {pkg_path.name}")
    if kind == "dlc" and not written:
        raise RuntimeError(f"no DLC EDAT found in {pkg_path.name}")

    return kind, title_id, folder