#!/usr/bin/env python3
"""
dokuro_uni_parser.py

Parser and extractor for Dokuro-chan PS2 UNI2 archives.

Known layout:
    Header      : 0x0000
    Entry table : 0x0800
    Sector size : 0x0800
    Data base   : header.data_start_sector * 0x800
    File offset : data base + entry.start_sector * 0x800

Usage:
    python dokuro_uni_parser.py script.uni
    python dokuro_uni_parser.py script.uni --extract extracted
    python dokuro_uni_parser.py script.uni --json metadata.json
"""

from __future__ import annotations

import argparse
import json
import struct
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import BinaryIO


UNI_SECTOR_SIZE = 0x800
UNI_TABLE_OFFSET = 0x800

HEADER_STRUCT = struct.Struct("<4sIIII")
ENTRY_STRUCT = struct.Struct("<IIII")

STCM2_PREFIX = b"STCM"
STCM2_SIGNATURE = b"STCM2 File Make By Minku 06.0"


class UniFormatError(Exception):
    """Raised when a UNI2 archive is malformed."""


@dataclass(slots=True)
class UniHeader:
    magic: str
    version: int
    entry_count: int
    flags: int
    data_start_sector: int

    @property
    def data_base(self) -> int:
        return self.data_start_sector * UNI_SECTOR_SIZE


@dataclass(slots=True)
class UniEntry:
    index: int
    file_id: int
    start_sector: int
    sector_count: int
    size: int
    absolute_offset: int
    allocated_size: int
    kind: str
    valid: bool
    warning: str = ""

    @property
    def end_offset(self) -> int:
        return self.absolute_offset + self.allocated_size


def read_exact(stream: BinaryIO, size: int) -> bytes:
    data = stream.read(size)
    if len(data) != size:
        raise UniFormatError(
            f"Unexpected end of file: wanted {size} bytes, got {len(data)}."
        )
    return data


def read_header(stream: BinaryIO, file_size: int) -> UniHeader:
    stream.seek(0)
    raw = read_exact(stream, HEADER_STRUCT.size)

    magic_raw, version, entry_count, flags, data_start_sector = (
        HEADER_STRUCT.unpack(raw)
    )

    magic = magic_raw.decode("ascii", errors="replace")

    if magic_raw != b"UNI2":
        raise UniFormatError(
            f"Invalid magic {magic_raw!r}; expected b'UNI2'."
        )

    if entry_count > 0x100000:
        raise UniFormatError(
            f"Unreasonable entry count: {entry_count}."
        )

    header = UniHeader(
        magic=magic,
        version=version,
        entry_count=entry_count,
        flags=flags,
        data_start_sector=data_start_sector,
    )

    if header.data_base > file_size:
        raise UniFormatError(
            f"Data base 0x{header.data_base:X} is outside the file "
            f"(size 0x{file_size:X})."
        )

    table_end = UNI_TABLE_OFFSET + entry_count * ENTRY_STRUCT.size
    if table_end > header.data_base:
        raise UniFormatError(
            f"Entry table ends at 0x{table_end:X}, overlapping data area "
            f"at 0x{header.data_base:X}."
        )

    if table_end > file_size:
        raise UniFormatError(
            f"Entry table extends beyond end of file."
        )

    return header


def detect_kind(stream: BinaryIO, offset: int, size: int) -> str:
    if size < 4:
        return "Raw data"

    stream.seek(offset)
    prefix = stream.read(min(size, len(STCM2_SIGNATURE)))

    if prefix.startswith(STCM2_SIGNATURE):
        return "STCM2 script"

    if prefix.startswith(STCM2_PREFIX):
        return "STCM-like data"

    return "Raw data"


def read_entries(
    stream: BinaryIO,
    header: UniHeader,
    file_size: int,
) -> list[UniEntry]:
    entries: list[UniEntry] = []

    stream.seek(UNI_TABLE_OFFSET)

    for index in range(header.entry_count):
        raw = read_exact(stream, ENTRY_STRUCT.size)
        file_id, start_sector, sector_count, size = ENTRY_STRUCT.unpack(raw)

        absolute_offset = (
            header.data_base + start_sector * UNI_SECTOR_SIZE
        )
        allocated_size = sector_count * UNI_SECTOR_SIZE

        warnings: list[str] = []

        if size > allocated_size:
            warnings.append(
                f"logical size 0x{size:X} exceeds allocation "
                f"0x{allocated_size:X}"
            )

        if absolute_offset > file_size:
            warnings.append(
                f"offset 0x{absolute_offset:X} is outside the file"
            )

        if absolute_offset + allocated_size > file_size:
            warnings.append(
                f"allocation ends at 0x{absolute_offset + allocated_size:X}, "
                f"past EOF 0x{file_size:X}"
            )

        valid = not warnings

        kind = "Invalid entry"
        if valid:
            kind = detect_kind(stream, absolute_offset, size)

        entries.append(
            UniEntry(
                index=index,
                file_id=file_id,
                start_sector=start_sector,
                sector_count=sector_count,
                size=size,
                absolute_offset=absolute_offset,
                allocated_size=allocated_size,
                kind=kind,
                valid=valid,
                warning="; ".join(warnings),
            )
        )

        # detect_kind() changes the file position, so restore the next
        # table entry offset explicitly.
        stream.seek(UNI_TABLE_OFFSET + (index + 1) * ENTRY_STRUCT.size)

    return entries


def print_header(header: UniHeader, file_size: int) -> None:
    print("UNI2 Header")
    print("-----------")
    print(f"Magic             : {header.magic}")
    print(f"Version           : 0x{header.version:08X}")
    print(f"Entry count       : {header.entry_count}")
    print(f"Flags             : 0x{header.flags:08X}")
    print(f"Data start sector : 0x{header.data_start_sector:X}")
    print(f"Data base         : 0x{header.data_base:X}")
    print(f"Archive size      : 0x{file_size:X} ({file_size} bytes)")
    print()


def print_entries(entries: list[UniEntry]) -> None:
    print(
        f"{'#':>3} "
        f"{'ID':>6} "
        f"{'StartSec':>10} "
        f"{'Sectors':>8} "
        f"{'Size':>10} "
        f"{'Offset':>10} "
        f"{'Allocated':>10} "
        f"{'Type'}"
    )
    print("-" * 96)

    for entry in entries:
        status = entry.kind
        if entry.warning:
            status = f"{status} | WARNING: {entry.warning}"

        print(
            f"{entry.index:3d} "
            f"{entry.file_id:6d} "
            f"0x{entry.start_sector:08X} "
            f"{entry.sector_count:8d} "
            f"0x{entry.size:08X} "
            f"0x{entry.absolute_offset:08X} "
            f"0x{entry.allocated_size:08X} "
            f"{status}"
        )


def choose_extension(entry: UniEntry) -> str:
    if entry.kind == "STCM2 script":
        return ".stcm2"
    return ".bin"


def extract_entries(
    stream: BinaryIO,
    entries: list[UniEntry],
    output_dir: Path,
    include_padding: bool,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    extracted = 0

    for entry in entries:
        if not entry.valid:
            print(
                f"Skipping invalid entry {entry.index} / ID {entry.file_id}: "
                f"{entry.warning}",
                file=sys.stderr,
            )
            continue

        read_size = entry.allocated_size if include_padding else entry.size

        stream.seek(entry.absolute_offset)
        data = read_exact(stream, read_size)

        extension = choose_extension(entry)
        output_name = (
            f"{entry.index:03d}_id_{entry.file_id}"
            f"_off_{entry.absolute_offset:08X}{extension}"
        )

        output_path = output_dir / output_name
        output_path.write_bytes(data)
        extracted += 1

    print()
    print(f"Extracted {extracted} entries to: {output_dir}")


def write_json(
    output_path: Path,
    source_path: Path,
    header: UniHeader,
    entries: list[UniEntry],
    file_size: int,
) -> None:
    document = {
        "source": str(source_path),
        "archive_size": file_size,
        "sector_size": UNI_SECTOR_SIZE,
        "table_offset": UNI_TABLE_OFFSET,
        "header": asdict(header),
        "entries": [asdict(entry) for entry in entries],
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(document, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )

    print()
    print(f"Wrote metadata to: {output_path}")


def build_argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Inspect and extract Dokuro-chan PS2 UNI2 archives."
        )
    )

    parser.add_argument(
        "archive",
        type=Path,
        help="Path to the UNI2 archive, for example script.uni.",
    )

    parser.add_argument(
        "--extract",
        metavar="DIRECTORY",
        type=Path,
        help="Extract all valid archive members.",
    )

    parser.add_argument(
        "--include-padding",
        action="store_true",
        help=(
            "When extracting, include the complete sector allocation instead "
            "of only the logical file size."
        ),
    )

    parser.add_argument(
        "--json",
        metavar="FILE",
        type=Path,
        help="Write parsed metadata to a JSON file.",
    )

    return parser


def main() -> int:
    parser = build_argument_parser()
    args = parser.parse_args()

    archive_path: Path = args.archive

    if not archive_path.is_file():
        print(
            f"ERROR: File not found: {archive_path}",
            file=sys.stderr,
        )
        return 1

    try:
        file_size = archive_path.stat().st_size

        with archive_path.open("rb") as stream:
            header = read_header(stream, file_size)
            entries = read_entries(stream, header, file_size)

            print_header(header, file_size)
            print_entries(entries)

            if args.extract is not None:
                extract_entries(
                    stream,
                    entries,
                    args.extract,
                    args.include_padding,
                )

            if args.json is not None:
                write_json(
                    args.json,
                    archive_path,
                    header,
                    entries,
                    file_size,
                )

    except (OSError, UniFormatError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1

    return 0


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