from os import path, makedirs
from json import loads, dump, dumps
from io import BytesIO
from logging import getLogger
from concurrent.futures import ProcessPoolExecutor
from time import sleep
from sssekai.unity.AssetBundle import load_assetbundle
from UnityPy.enums import ClassIDType
from sssekai.fmt.rla import read_rla, read_rla_frame
from tqdm import tqdm
import zipfile, base64

logger = getLogger(__name__)


def dump_json_job(sname, version, script):
    rla = read_rla(BytesIO(script), version)
    dump(
        rla, open(sname + ".json", "w", encoding="utf-8"), indent=4, ensure_ascii=False
    )


def dump_audio_job(sname, version, script):
    rla = read_rla(BytesIO(script), version)
    for tick, data in rla.items():
        data = data.get("SoundData", None)
        if data:
            for i, b64data in enumerate(data):
                fname = sname + "_%d_%d.hca" % (i, tick)
                raw_data = base64.b64decode(b64data["data"])
                start = raw_data.find(b"\xFF\xFF")
                with open(fname, "wb") as f:
                    f.write(raw_data[start:])


def main_rla2json(args):
    with open(args.infile, "rb") as f:
        datas = dict()
        if f.read(2) == b"PK":
            f.seek(0)
            with zipfile.ZipFile(f, "r") as z:
                for name in z.namelist():
                    with z.open(name) as zf:
                        datas[name] = zf.read()
        f.seek(0)
        if f.read(4) == b"RTVL":
            f.seek(0)
            frame = read_rla_frame(
                f.read(), version=tuple(map(int, args.version.split(".")))
            )
            print(dumps(frame, indent=4, ensure_ascii=False))
            return
        else:
            f.seek(0)
            rla_env = load_assetbundle(f)
            for obj in rla_env.objects:
                if obj.type in {ClassIDType.TextAsset}:
                    data = obj.read()
                    datas[data.m_Name] = data.m_Script.encode(
                        "utf-8", "surrogateescape"
                    )
        header = datas.get("sekai.rlh", None)
        assert header, "RLH Header file not found!"
        makedirs(args.outdir, exist_ok=True)
        header = loads(header.decode("utf-8"))
        splitSeconds = header["splitSeconds"]

        dump(
            header,
            open(path.join(args.outdir, "sekai.rlh.json"), "w", encoding="utf-8"),
            indent=4,
            ensure_ascii=False,
        )
        version = tuple(map(int, header["version"].split(".")))
        logger.info("Version: %d.%d" % version)
        logger.info("Count: %d" % len(header["splitFileIds"]))

        worker_job = dump_json_job
        if args.dump_audio:
            worker_job = dump_audio_job
            logger.info("Dumping RLA raw HCA audio data")
        else:
            logger.info("Dumping RLA data to JSON")
        with ProcessPoolExecutor() as executor:
            if not args.no_parallel:
                logger.info(
                    "Dumping RLA data with %d processors" % executor._max_workers
                )
            futures = []
            for sid in header["splitFileIds"]:
                sname = "sekai_%02d_%08d" % (splitSeconds, sid)
                script = datas[sname + ".rla"]
                if args.no_parallel:
                    worker_job(path.join(args.outdir, sname), version, script)
                else:
                    futures.append(
                        executor.submit(
                            worker_job, path.join(args.outdir, sname), version, script
                        )
                    )
            finsihed_futures = set()
            with tqdm(total=len(futures)) as pbar:
                while len(finsihed_futures) < len(futures):
                    for i, future in enumerate(futures):
                        if future.done() and i not in finsihed_futures:
                            pbar.update(1)
                            finsihed_futures.add(i)
                    sleep(0.1)
