# pyright: reportAny=false
"""
Script to extract arc.pack files used by Arcaea Nintendo Switch.

Usage:

- Extract the contents of an arc.pack file:
    python arcpack.py x arc.pack arc.json output
"""
import argparse
import json
import os
from pathlib import Path
from typing import TypedDict, cast


class ArcJSONEntry(TypedDict):
    OriginalFilename: str
    Offset: int
    Length: int


class ArcJSONGroup(TypedDict):
    Name: str
    Offset: int
    Length: int
    OrderedEntries: list[ArcJSONEntry]


class ArcJSON(TypedDict):
    Groups: list[ArcJSONGroup]


def arcpack_extract(args: argparse.Namespace):
    arcpack = cast(Path, args.arcpack)
    arcjson = cast(Path, args.arcjson)
    output = cast(Path, args.output_directory)

    with arcjson.open("r", encoding="utf-8") as f:
        toc: ArcJSON = json.load(f)
    
    try:
        from tqdm import tqdm

        total = 0

        for group in toc["Groups"]:
            total += len(group["OrderedEntries"])

        pb = tqdm(
            desc="Extracting pack",
            total=total,
        )
    except ImportError:
        pb = None
    
    with arcpack.open("rb") as f:
        for group in toc["Groups"]:
            for entry in group["OrderedEntries"]:
                target = output / entry["OriginalFilename"]

                target.mkdir(parents=True, exist_ok=True)

                _ = f.seek(entry["Offset"], os.SEEK_SET)

                with target.open("wb") as ft:
                    _ = ft.write(f.read(entry["Length"]))

                if pb is not None:
                    _ = pb.update(1)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="arc.pack extractor")
    subcommands = parser.add_subparsers(required=True)

    extract_command = subcommands.add_parser("x", help="Extract")
    _ = extract_command.add_argument("arcpack", type=Path)
    _ = extract_command.add_argument("arcjson", type=Path)
    _ = extract_command.add_argument("output_directory", type=Path)

    args = parser.parse_args()
    args.func(args)