import io
import os
import UnityPy
from UnityPy import enums
from UnityPy.files import ObjectReader
from PIL import Image

def convert(input_path : str, output_path : str, platform = enums.BuildTarget.WebGL, enable_read = False):
    env = UnityPy.load(input_path)

    for item in env.assets:
        item._m_target_platform = int(platform)

    for obj in env.objects:
        if obj.type == enums.ClassIDType.Texture2D:
            convert_texture(obj, platform)
        if enable_read and obj.type == enums.ClassIDType.Mesh:
            data = obj.read()
            data.m_IsReadable = True
            data.save()

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, "wb") as f:
        f.write(env.file.save(packer="lz4"))

def extract(input_path : str, output_path : str):
    env = UnityPy.load(input_path)

    for obj in env.objects:
        if obj.type == enums.ClassIDType.Texture2D:
            os.makedirs(output_path, exist_ok=True)
            data = obj.read()
            data.image.save(output_path + data.name + ".png", format="PNG")

def convert_texture(obj : ObjectReader, target_platform : enums.BuildTarget):
    if target_platform == enums.BuildTarget.WebGL:
        data = obj.read()
        with io.BytesIO() as output:
            data.image.save(output, format="PNG")
            if data.image.width % 4 == 0 and data.image.height % 4 == 0:
                data.m_TextureFormat = enums.TextureFormat.DXT5
            else:
                data.m_TextureFormat = enums.TextureFormat.ARGB32
            data.image = Image.open(output)
        data.save()