#!/usr/bin/env python3
import argparse
import os
import re
import sys
from collections import defaultdict
from collections.abc import Iterable
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from shared.paths import DATA_DIR, DOCS_DIR, REPO_DIR
from shared.vfs import VirtualFilesystem

# The install-tree docs list the images that ship from the TRX-data repo. Point
# TRX_DATA_DIR at a local checkout to regenerate them; without it the tool has
# nothing to read and bows out.
DATA_REPO_ROOT_DIR = Path(
    os.environ.get("TRX_DATA_DIR", REPO_DIR.parent / "TRX-data")
)

IGNORED_FILES = """
max_stats.cache.json
"""

TR1_OG_FILES = """
data/cut1.phd
data/cut2.phd
data/cut3.phd
data/cut4.phd
data/gym.phd
data/level1.phd
data/level2.phd
data/level3a.phd
data/level3b.phd
data/level4.phd
data/level5.phd
data/level6.phd
data/level7a.phd
data/level7b.phd
data/level8a.phd
data/level8b.phd
data/level8c.phd
data/level10a.phd
data/level10b.phd
data/level10c.phd
data/title.phd
fmv/cafe.rpl
fmv/canyon.rpl
fmv/core.avi
fmv/end.rpl
fmv/escape.rpl
fmv/lift.rpl
fmv/mansion.rpl
fmv/prison.rpl
fmv/pyramid.rpl
fmv/snow.rpl
fmv/vision.rpl
music/track02.flac
music/track03.flac
music/track04.flac
music/track05.flac
music/track06.flac
music/track07.flac
music/track08.flac
music/track09.flac
music/track10.flac
music/track11.flac
music/track12.flac
music/track13.flac
music/track14.flac
music/track15.flac
music/track16.flac
music/track17.flac
music/track18.flac
music/track19.flac
music/track20.flac
music/track21.flac
music/track22.flac
music/track23.flac
music/track24.flac
music/track25.flac
music/track26.flac
music/track27.flac
music/track28.flac
music/track29.flac
music/track30.flac
music/track31.flac
music/track32.flac
music/track33.flac
music/track34.flac
music/track35.flac
music/track36.flac
music/track37.flac
music/track38.flac
music/track39.flac
music/track40.flac
music/track41.flac
music/track42.flac
music/track43.flac
music/track44.flac
music/track45.flac
music/track46.flac
music/track47.flac
music/track48.flac
music/track49.flac
music/track50.flac
music/track51.flac
music/track52.flac
music/track53.flac
music/track54.flac
music/track55.flac
music/track56.flac
music/track57.flac
music/track58.flac
music/track59.flac
music/track60.flac
"""

TR1_UB_OG_FILES = """
data/cat.phd
data/egypt.phd
data/end.phd
data/end2.phd
"""

TR2_OG_FILES = """
data/assault.tr2
data/boat.tr2
data/catacomb.tr2
data/cut1.tr2
data/cut2.tr2
data/cut3.tr2
data/cut4.tr2
data/deck.tr2
data/emprtomb.tr2
data/floating.tr2
data/house.tr2
data/icecave.tr2
data/keel.tr2
data/living.tr2
data/main.sfx
data/monastry.tr2
data/opera.tr2
data/platform.tr2
data/rig.tr2
data/skidoo.tr2
data/title.tr2
data/unwater.tr2
data/venice.tr2
data/wall.tr2
data/xian.tr2
fmv/ancient.rpl
fmv/crash.rpl
fmv/end.rpl
fmv/jeep.rpl
fmv/landing.rpl
fmv/logo.rpl
fmv/modern.rpl
fmv/ms.rpl
music/2.mp3
music/3.mp3
music/4.mp3
music/5.mp3
music/6.mp3
music/7.mp3
music/8.mp3
music/9.mp3
music/10.mp3
music/11.mp3
music/12.mp3
music/13.mp3
music/14.mp3
music/15.mp3
music/16.mp3
music/17.mp3
music/18.mp3
music/19.mp3
music/20.mp3
music/21.mp3
music/22.mp3
music/23.mp3
music/24.mp3
music/25.mp3
music/26.mp3
music/27.mp3
music/28.mp3
music/29.mp3
music/30.mp3
music/31.mp3
music/32.mp3
music/33.mp3
music/34.mp3
music/35.mp3
music/36.mp3
music/37.mp3
music/38.mp3
music/39.mp3
music/40.mp3
music/41.mp3
music/42.mp3
music/43.mp3
music/44.mp3
music/45.mp3
music/46.mp3
music/47.mp3
music/48.mp3
music/49.mp3
music/50.mp3
music/51.mp3
music/52.mp3
music/53.mp3
music/54.mp3
music/55.mp3
music/56.mp3
music/57.mp3
music/58.mp3
music/59.mp3
music/60.mp3
music/61.mp3
"""

TR2_GM_OG_FILES = """
data/level1.tr2
data/level2.tr2
data/level3.tr2
data/level4.tr2
data/level5.tr2
data/main_gm.sfx
data/title_gm.tr2
"""

TR3_OG_FILES = """
data/main.sfx
data/tombpc.dat
fmv/crsh_eng.rpl
fmv/endgame.rpl
fmv/intr_eng.rpl
fmv/sail_eng.rpl
fmv/logo.rpl
audio/cdaudio.wad
cuts/cut1.tr2
cuts/cut2.tr2
data/antarc.tr2
data/chamber.tr2
data/city.tr2
data/compound.tr2
data/crash.tr2
data/house.tr2
data/jungle.tr2
data/mines.tr2
data/nevada.tr2
data/office.tr2
data/quadchas.tr2
data/rapids.tr2
data/roofs.tr2
data/sewer.tr2
data/shore.tr2
data/stpaul.tr2
data/temple.tr2
data/title.tr2
data/tonyboss.tr2
data/tower.tr2
data/triboss.tr2
cuts/cut3.tr2
cuts/cut4.tr2
cuts/cut5.tr2
cuts/cut6.tr2
cuts/cut7.tr2
cuts/cut8.tr2
cuts/cut9.tr2
cuts/cut11.tr2
cuts/cut12.tr2
data/area51.tr2
"""

TR3_LA_OG_FILES = """
data/chunnel.tr2
data/scotland.tr2
data/slinc.tr2
data/undersea.tr2
data/willsden.tr2
data/zoo.tr2
data/main_la.sfx
data/title_la.tr2
"""

OG_FILE_GROUPS = {
    "tr1": tuple(Path(line) for line in TR1_OG_FILES.splitlines() if line),
    "tr1-ub": tuple(Path(line) for line in TR1_UB_OG_FILES.splitlines() if line),
    "tr2": tuple(Path(line) for line in TR2_OG_FILES.splitlines() if line),
    "tr2-gm": tuple(Path(line) for line in TR2_GM_OG_FILES.splitlines() if line),
    "tr3": tuple(Path(line) for line in TR3_OG_FILES.splitlines() if line),
    "tr3-la": tuple(Path(line) for line in TR3_LA_OG_FILES.splitlines() if line),
}

OG_PREFIX_TARGETS: dict[str, dict[str, Path]] = {
    "tr1": {
        "data": Path("games/tr1/levels"),
        "fmv": Path("games/tr1/fmv"),
        "music": Path("games/tr1/music"),
    },
    "tr1-ub": {
        "data": Path("games/tr1-ub/levels"),
    },
    "tr2": {
        "data": Path("games/tr2/levels"),
        "fmv": Path("games/tr2/fmv"),
        "music": Path("games/tr2/music"),
    },
    "tr2-gm": {
        "data": Path("games/tr2-gm/levels"),
    },
    "tr3": {
        "data": Path("games/tr3/levels"),
        "fmv": Path("games/tr3/fmv"),
        "audio": Path("games/tr3/audio"),
        "cuts": Path("games/tr3/cuts"),
    },
    "tr3-la": {
        "data": Path("games/tr3-la/levels"),
    },
}

OG_FILE_OVERRIDES: dict[str, dict[Path, Path]] = {
    "tr2": {
        Path("data/main.sfx"): Path("games/tr2/main.sfx"),
    },
    "tr2-gm": {
        Path("data/main_gm.sfx"): Path("games/tr2-gm/main.sfx"),
        Path("data/title_gm.tr2"): Path("games/tr2-gm/levels/title.tr2"),
    },
    "tr3": {
        Path("data/main.sfx"): Path("games/tr3/main.sfx"),
        Path("data/tombpc.dat"): Path("games/tr3/tombpc.dat"),
    },
    "tr3-la": {
        Path("data/main_la.sfx"): Path("games/tr3-la/main.sfx"),
        Path("data/title_la.tr2"): Path("games/tr3-la/levels/title.tr2"),
    },
}

TRX_IMAGE_ROOTS = {
    "tr1": DATA_REPO_ROOT_DIR / "tr1/ship/data/images",
    "tr1-ub": DATA_REPO_ROOT_DIR / "tr1-ub/ship/data/images",
    "tr2": DATA_REPO_ROOT_DIR / "tr2/ship/data/images",
    "tr2-gm": DATA_REPO_ROOT_DIR / "tr2-gm/ship/data/images",
    "tr3": DATA_REPO_ROOT_DIR / "tr3/ship/data/images",
    "tr3-la": DATA_REPO_ROOT_DIR / "tr3-la/ship/data/images",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Update game strings JSON files to reflect GS_DEFINE() macros."
    )
    parser.add_argument(
        "-d",
        "--dry-run",
        action="store_true",
        help="Perform a dry run: show unified diffs instead of writing files.",
    )
    return parser.parse_args()


def collect(root_dir: Path) -> Iterable[Path]:
    paths = [
        path.relative_to(root_dir)
        for path in root_dir.rglob("*")
        if path.is_file() and path.name not in IGNORED_FILES
    ]
    return sorted(paths)


def remap_og_path(path: Path, game: str) -> Path | None:
    override = OG_FILE_OVERRIDES.get(game, {}).get(path)
    if override is not None:
        return override
    if path.parts[0] in OG_PREFIX_TARGETS[game]:
        return OG_PREFIX_TARGETS[game][path.parts[0]] / path.relative_to(
            path.parts[0]
        )
    return None


def remap_og_paths(paths: Iterable[Path], game: str) -> Iterable[Path]:
    for path in paths:
        mapped = remap_og_path(path, game)
        if mapped is not None:
            yield mapped


def iter_trx_image_paths() -> Iterable[Path]:
    for game, root_dir in TRX_IMAGE_ROOTS.items():
        for path in collect(root_dir):
            yield Path("games") / game / "images" / path


def build_trx_win_paths() -> Iterable[Path]:
    yield Path("TRX.exe")
    yield Path("cfg/shell.json5*")
    yield Path("cfg/TR1X.json5*")
    yield Path("cfg/TR2X.json5*")
    yield Path("cfg/TR3X.json5*")
    yield from collect(DATA_DIR / "trx/ship")
    yield from iter_trx_image_paths()
    for group_name in OG_FILE_GROUPS:
        yield from remap_og_paths(OG_FILE_GROUPS[group_name], group_name)


def build_trx_mac_paths() -> Iterable[Path]:
    yield Path("Contents/_CodeSignature")
    yield Path("Contents/Frameworks")
    yield Path("Contents/Info.plist")
    yield Path("Contents/MacOS")
    yield from (
        Path("Contents/Resources") / sub
        for sub in build_trx_win_paths()
        if sub != Path("TRX.exe")
    )
    yield Path("Contents/Resources/icon.icns")


def build_doc_trees() -> Iterable[tuple[Path, str, list[Path]]]:
    yield (DOCS_DIR / "trx/INSTALLING.md", "win", list(build_trx_win_paths()))
    yield (DOCS_DIR / "trx/INSTALLING.md", "mac", list(build_trx_mac_paths()))




def make_tree_string(paths: list[Path]) -> str:
    type Node = defaultdict[Any, Any]

    # build nested dict
    def tree() -> Node:
        return defaultdict(tree)

    root = tree()
    for path in paths:
        node = root
        for part in str(path).split("/"):
            node = node[part]

    def natural_sort_key(s: str) -> list[str | int]:
        return [
            int(text) if text.isdigit() else text.lower()
            for text in re.split("([0-9]+)", s)
        ]

    # recursive render
    def render(node: Node, prefix: str = "") -> list[str]:
        entries = sorted(
            node.keys(),
            key=lambda name: (0 if node[name] else 1, natural_sort_key(name)),
        )
        lines = []
        for i, name in enumerate(entries):
            last = i == len(entries) - 1
            branch = "└── " if last else "├── "
            subprefix = "    " if last else "│   "
            lines.append(prefix + branch + name)
            if node[name]:
                lines.extend(render(node[name], prefix + subprefix))
        return lines

    return ".\n" + "\n".join(render(root))


def process_install_doc(content: str, platform: str, paths: list[Path]) -> str:
    result: list[str] = []
    lines = content.splitlines()
    while lines:
        line = lines.pop(0)
        if line == f'<details data-id="file-tree-{platform}">':
            result.append(line)
            while lines:
                line = lines.pop(0)
                if not line:
                    break
            result.append(
                "<pre><code>" + make_tree_string(paths) + "</code></pre>"
            )
            result.append("</details>")
            if lines:
                result.append(line)
        else:
            result.append(line)
    return "\n".join(result) + "\n"


def main() -> None:
    args = parse_args()

    if not DATA_REPO_ROOT_DIR.exists():
        print(
            f"install_trees: TRX-data not found at {DATA_REPO_ROOT_DIR}, "
            "skipping",
            file=sys.stderr,
        )
        return

    vfs = VirtualFilesystem()

    for path, platform, paths in build_doc_trees():
        vfs.put(
            path, process_install_doc(vfs.get(path), platform, paths)
        )

    if args.dry_run:
        vfs.show_diff()
    else:
        vfs.commit()


if __name__ == "__main__":
    main()
