import argparse
import logging
import re
import sys
import zipfile
from pathlib import Path, PurePath

from tqdm import tqdm

from common import ArcaeaApkParser, ExtractTask, extract

ROOT_OUTPUT_PATH = Path(sys.argv[0]).parent.absolute()

logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser(
    prog="arcaea-apk-extract for Andreal",
    description="Literal meaning.",
    epilog="This program is licensed under the WTFPL. Feel free to fuck anything.",
)
parser.add_argument("apk_file")
parser.add_argument(
    "-op",
    "--output-path",
    action="store",
    help="output path, defaults to Path(sys.argv[0]).parent",
)
parser.add_argument(
    "-v",
    "--verbose",
    action="store_true",
    help="set logging level to logging.DEBUG",
)


class SongParser(ArcaeaApkParser):
    @property
    def OUTPUT_PATH(self):
        return ROOT_OUTPUT_PATH / "Song"

    def parse(self):
        songs_dir_zf_path = zipfile.Path(self.zf) / "assets" / "songs"
        songs_zf_path = [zfp for zfp in songs_dir_zf_path.iterdir() if zfp.is_dir()]

        tasks = []
        for song_zf_path in songs_zf_path:
            for file_zf_path in song_zf_path.iterdir():
                if not file_zf_path.is_file():
                    continue

                file_pure_path = PurePath(str(file_zf_path))

                if file_pure_path.suffix not in [".jpg", ".png"]:
                    continue

                song_id = re.sub(r"^dl_", "", song_zf_path.name)
                if file_pure_path.stem in ["base", "1080_base"]:
                    new_file_name_stem = song_id
                elif file_pure_path.stem in ["1080_base_night", "base_night"]:
                    new_file_name_stem = f"{song_id}_night"
                elif file_pure_path.stem in [
                    "0",
                    "1",
                    "2",
                    "3",
                    "1080_0",
                    "1080_1",
                    "1080_2",
                    "1080_3",
                ]:
                    new_file_name_stem = (
                        f"{song_id}_{file_pure_path.stem.replace('1080_', '')}"
                    )
                else:
                    continue

                new_file_name = f"{new_file_name_stem}{file_pure_path.suffix}"
                new_file_path = self.OUTPUT_PATH / new_file_name

                tasks.append(ExtractTask(file_zf_path, new_file_path))
        return tasks


class CharParser(ArcaeaApkParser):
    CHAR_RE = r"^\d+u?\.png$"
    CHAR_ICON_RE = r"^\d+u?_icon\.png$"

    @property
    def OUTPUT_PATH_CHAR(self):
        return ROOT_OUTPUT_PATH / "Char"

    @property
    def OUTPUT_PATH_ICON(self):
        return ROOT_OUTPUT_PATH / "Icon"

    def parse(self):
        if (zipfile.Path(self.zf) / "assets" / "char" / "1080").exists():
            char_dir_zf_path = zipfile.Path(self.zf) / "assets" / "char" / "1080"
        else:
            char_dir_zf_path = zipfile.Path(self.zf) / "assets" / "char"
        char_zf_path_list = [zfp for zfp in char_dir_zf_path.iterdir() if zfp.is_file()]
        char_icon_dir_zf_path = zipfile.Path(self.zf) / "assets" / "char"
        char_icon_zf_path_list = [
            zfp for zfp in char_icon_dir_zf_path.iterdir() if zfp.is_file()
        ]

        tasks = []

        for char_zf_path in char_zf_path_list:
            char_pure_path = PurePath(str(char_zf_path))
            if not re.match(self.CHAR_RE, char_pure_path.name):
                continue

            new_file_name = char_pure_path.name
            new_file_path = self.OUTPUT_PATH_CHAR / new_file_name

            tasks.append(ExtractTask(char_zf_path, new_file_path))

        for char_icon_zf_path in char_icon_zf_path_list:
            char_icon_pure_path = PurePath(str(char_icon_zf_path))
            if not re.match(self.CHAR_ICON_RE, char_icon_pure_path.name):
                continue

            new_file_name = char_icon_pure_path.name.replace("_icon", "")
            new_file_path = self.OUTPUT_PATH_ICON / new_file_name

            tasks.append(ExtractTask(char_icon_zf_path, new_file_path))
        return tasks


if __name__ == "__main__":
    args = parser.parse_args(sys.argv[1:])

    if args.output_path:
        output_path = Path(args.output_path)
        assert output_path.exists()
        ROOT_OUTPUT_PATH = output_path

    logging_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(
        level=logging_level,
        style="{",
        format="[{levelname}]{name}: {msg}",
        stream=sys.stdout,
        encoding="utf-8",
    )

    apk_path = Path(args.apk_file)
    assert apk_path.exists()
    apk_zf = zipfile.ZipFile(str(apk_path))

    sp = SongParser(apk_zf)
    cp = CharParser(apk_zf)

    sp.OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
    cp.OUTPUT_PATH_CHAR.mkdir(parents=True, exist_ok=True)
    cp.OUTPUT_PATH_ICON.mkdir(parents=True, exist_ok=True)

    sp_tasks = sp.parse()
    logger.info(f"{len(sp_tasks)} files from songs")
    cp_tasks = cp.parse()
    logger.info(f"{len(cp_tasks)} files from char")

    tasks = sp_tasks + cp_tasks
    logger.info("Extracting, please wait...")
    for task in tqdm(tasks):
        extract(task)