#!/usr/bin/env python3

"""Add native solo RDF categories to a WotLK 3.3.5a LFGDungeons.dbc."""

from __future__ import annotations

import argparse
import struct
from pathlib import Path


MAGIC = b"WDBC"
FIELD_COUNT = 49
RECORD_SIZE = FIELD_COUNT * 4
SOLO_ID_OFFSET = 1000

SOLO_NAMES = {
    258: "Случайное одиночное подземелье: Классика",
    259: "Случайное одиночное подземелье: Burning Crusade",
    260: "Случайное одиночное героическое: Burning Crusade",
    261: "Случайное одиночное подземелье Lich King",
    262: "Случайное одиночное героическое подземелье Lich King",
}

SOLO_DESCRIPTION = (
    "Случайное подземелье для одного игрока. "
    "Существа и боссы масштабируются системой AutoBalance."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Duplicate the five stock random dungeon rows as solo RDF categories."
    )
    parser.add_argument("source", type=Path, help="Source LFGDungeons.dbc")
    parser.add_argument("output", type=Path, help="Patched LFGDungeons.dbc")
    return parser.parse_args()


def add_string(string_block: bytearray, value: str) -> int:
    encoded = value.encode("utf-8") + b"\0"
    offset = len(string_block)
    string_block.extend(encoded)
    return offset


def main() -> None:
    args = parse_args()
    raw = args.source.read_bytes()
    if len(raw) < 20 or raw[:4] != MAGIC:
        raise SystemExit(f"{args.source}: not a WDBC file")

    record_count, field_count, record_size, string_size = struct.unpack_from("<4I", raw, 4)
    if field_count != FIELD_COUNT or record_size != RECORD_SIZE:
        raise SystemExit(
            f"{args.source}: unexpected layout "
            f"(fields={field_count}, record_size={record_size})"
        )

    records_start = 20
    strings_start = records_start + record_count * record_size
    strings_end = strings_start + string_size
    if strings_end != len(raw):
        raise SystemExit(f"{args.source}: inconsistent WDBC size")

    records: list[list[int]] = []
    source_records: dict[int, list[int]] = {}
    solo_ids = {base_id + SOLO_ID_OFFSET for base_id in SOLO_NAMES}

    for index in range(record_count):
        offset = records_start + index * record_size
        record = list(struct.unpack_from(f"<{FIELD_COUNT}I", raw, offset))
        if record[0] in SOLO_NAMES:
            source_records[record[0]] = record
        if record[0] not in solo_ids:
            records.append(record)

    missing = sorted(set(SOLO_NAMES) - set(source_records))
    if missing:
        raise SystemExit(f"{args.source}: missing stock random dungeon rows: {missing}")

    string_block = bytearray(raw[strings_start:strings_end])
    for base_id, name in SOLO_NAMES.items():
        record = source_records[base_id].copy()
        record[0] = base_id + SOLO_ID_OFFSET

        name_offset = add_string(string_block, name)
        description_offset = add_string(string_block, SOLO_DESCRIPTION)

        # Fill every locale slot. This keeps the custom row readable regardless
        # of which locale column the client selects.
        record[1:17] = [name_offset] * 16
        record[32:48] = [description_offset] * 16
        records.append(record)

    records.sort(key=lambda record: record[0])
    output = bytearray()
    output.extend(MAGIC)
    output.extend(struct.pack("<4I", len(records), FIELD_COUNT, RECORD_SIZE, len(string_block)))
    for record in records:
        output.extend(struct.pack(f"<{FIELD_COUNT}I", *record))
    output.extend(string_block)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
    temporary.write_bytes(output)
    temporary.replace(args.output)
    print(f"Wrote {args.output} with solo RDF entries 1258-1262")


if __name__ == "__main__":
    main()