import os
from sssekai.unity.AssetBundle import load_assetbundle
from UnityPy.enums import ClassIDType
from logging import getLogger

logger = getLogger(__name__)


def main_spineextract(args):
    outdir = args.outdir
    os.makedirs(outdir, exist_ok=True)
    with open(args.infile, "rb") as f:
        env = load_assetbundle(f)
        objects = [pobj.read() for pobj in env.objects]
        binaries = {
            obj.name: obj
            for obj in objects
            if getattr(obj, "type", None) in {ClassIDType.TextAsset}
        }
        textures = {
            obj.name: obj
            for obj in objects
            if getattr(obj, "type", None) in {ClassIDType.Texture2D}
        }
        spines = set()
        for name in binaries:
            if name.endswith(".atlas") or name.endswith(".skel"):
                spines.add(".".join(name.split(".")[:-1]))
        for spine in spines:
            logger.info("Extracting %s" % spine)
            atlas = binaries.get(spine + ".atlas", None)
            skel = binaries.get(spine + ".skel", None)
            os.makedirs(os.path.join(outdir, spine), exist_ok=True)
            if atlas:
                logger.info("...has Atlas %s" % spine)
                with open(os.path.join(outdir, spine, spine + ".atlas.txt"), "wb") as f:
                    f.write(atlas.script)
                texfiles = [line.strip() for line in atlas.text.split("\n")]
                texfiles = [
                    ".".join(line.split(".")[:-1])
                    for line in texfiles
                    if (line.endswith(".png"))
                ]
                for tex in texfiles:
                    logger.info("...has Texture %s" % tex)
                    textureobj = textures.get(tex, None)
                    if textureobj:
                        textureobj.image.save(os.path.join(outdir, spine, tex + ".png"))
                    else:
                        logger.warning("No texture found for %s" % tex)
            else:
                logger.warning(
                    "No atlas found for %s. Consequnetially, no textures cannot be exported either."
                    % spine
                )

            if skel:
                logger.info("...has Skeleton %s" % spine)
                with open(
                    os.path.join(outdir, spine, spine + ".skel.bytes"), "wb"
                ) as f:
                    f.write(skel.script)
            else:
                logger.warning("No skel found for %s" % spine)
