"""Discover type-0x01 secondary stems/keys from helper + primary type-0x07."""

from __future__ import annotations

import hashlib
import re
import struct
from pathlib import Path

from _secondary_pack import F00D_MAGIC, chacha_crypt, decrypt_secondary_minjs, unwrap_xz

NAME_RE = re.compile(rb"([0-9a-f]{40}\.min\.js)")
DEFAULT_HELPER_STEM = "7a7d99099b035b2c6512b6ebeeea6df1ede70fbb"


def parse_f00d_entries(data: bytes) -> list[dict]:
    magic, count = struct.unpack_from("<II", data, 0)
    if magic != F00D_MAGIC:
        raise ValueError("not F00DBEEF")
    if not 1 <= count <= 16:
        raise ValueError(f"unexpected F00DBEEF entry count {count}")
    entries = []
    for i in range(count):
        typ, flags, offset, size = struct.unpack_from("<4I", data, 8 + i * 16)
        payload = data[offset : offset + size]
        if len(payload) != size:
            raise ValueError("truncated F00DBEEF payload")
        entries.append(
            {
                "type": (typ >> 16) & 0xFFFF,
                "type_raw": typ,
                "flags": flags,
                "offset": offset,
                "size": size,
                "payload": payload,
            }
        )
    return entries


def parse_07_slots(payload: bytes) -> list[dict]:
    """type-0x07 records: 32-byte ChaCha key immediately before `{40hex}.min.js`."""
    slots = []
    for match in NAME_RE.finditer(payload):
        name = match.group(1).decode("ascii")
        start = match.start()
        if start < 32:
            continue
        key = payload[start - 32 : start]
        slots.append({"stem": name[:-7], "file": name, "key": key.hex()})
    return slots


def decrypt_pack(blob: bytes, key: bytes) -> bytes:
    return unwrap_xz(chacha_crypt(blob, key))


def discover_secondary_stems(
    weifile: Path,
    helper_key: bytes,
    helper_stem: str = DEFAULT_HELPER_STEM,
) -> dict[str, dict]:
    """Return stem -> {key, size, sha256} for every type-0x01 pointed at by primaries."""
    helper_path = weifile / f"{helper_stem}.min.js"
    if not helper_path.is_file():
        raise FileNotFoundError(f"missing helper pack: {helper_path}")
    helper_plain = decrypt_pack(helper_path.read_bytes(), helper_key)
    primary_slots = []
    for entry in parse_f00d_entries(helper_plain):
        if entry["type"] == 7:
            primary_slots.extend(parse_07_slots(entry["payload"]))
    if not primary_slots:
        raise ValueError("helper type-0x07 had no primary slots")

    stems: dict[str, dict] = {}
    for slot in primary_slots:
        primary_path = weifile / slot["file"]
        if not primary_path.is_file():
            raise FileNotFoundError(f"missing primary pack: {primary_path}")
        primary_plain = decrypt_pack(primary_path.read_bytes(), bytes.fromhex(slot["key"]))
        for entry in parse_f00d_entries(primary_plain):
            if entry["type"] != 7:
                continue
            for secondary in parse_07_slots(entry["payload"]):
                stem = secondary["stem"]
                path = weifile / secondary["file"]
                if not path.is_file():
                    raise FileNotFoundError(f"missing secondary pack: {path}")
                dylib = decrypt_secondary_minjs(path.read_bytes(), bytes.fromhex(secondary["key"]))
                digest = hashlib.sha256(dylib).hexdigest()
                prev = stems.get(stem)
                if prev and prev["key"] != secondary["key"]:
                    raise ValueError(f"conflicting keys for {stem}")
                stems[stem] = {
                    "key": secondary["key"],
                    "size": len(dylib),
                    "sha256": digest,
                    "dylib": dylib,
                }
    if not stems:
        raise ValueError("no type-0x01 secondaries discovered")
    return stems


def group_by_dylib_hash(stems: dict[str, dict]) -> dict[str, bytes]:
    """Unique plaintext dylibs keyed by sha256."""
    groups: dict[str, bytes] = {}
    for info in stems.values():
        digest = info["sha256"]
        if digest not in groups:
            groups[digest] = info["dylib"]
        elif groups[digest] != info["dylib"]:
            raise ValueError(f"dylib hash collision for {digest}")
    return groups