import os
import numpy as np
import UnityPy
from PIL import Image

#made by Katboi01
def convert(in_path, out_path, file_name):
    env = UnityPy.load(os.path.join(in_path, file_name))

    textures = []
    for obj in env.objects:
        if obj.type.name == "Texture2D":
            textures.append(obj)

    if len(textures) > 1:
        out_path = os.path.join(out_path, file_name)
        os.makedirs(out_path, exist_ok=True)

    for idx, obj in enumerate(textures):
        try:
            data = obj.read()
            im : Image = data.image
            width, height = im.size

            left = int(width * 0.8)
            padding = 1

            im1 = im.crop((0, 0, left - padding, height))
            im1_array = np.asarray(im1, dtype=np.float32) / 255
            
            y_channel = im1_array[:, :, 0]
            alpha_channel = im1_array[:, :, 2]

            im2 = im.crop((left + padding, 0, width, height)).resize((left - padding, height), Image.NEAREST)
            im2_array = np.asarray(im2, dtype=np.float32) / 255

            cr_channel = im2_array[:, :, 0]
            cb_channel = im2_array[:, :, 2]

            cr_channel -= 0.5
            cb_channel -= 0.5

            r_channel = np.clip(y_channel + (1.402 * cr_channel), 0, 1)
            g_channel = np.clip(y_channel + (-0.344136 * cb_channel) + (-0.714136 * cr_channel), 0, 1)
            b_channel = np.clip(y_channel + (1.772 * cb_channel), 0, 1)

            new_image = np.stack([r_channel, g_channel, b_channel, alpha_channel], axis=-1)
            new_image = (new_image * 255).astype(np.uint8)

            result = Image.fromarray(new_image)

            result.save(os.path.join(out_path, data.m_Name + ".png"), format="PNG")
        except Exception as e:
            print(f"Failed to extract ({idx}/{len(textures)}) from {file_name}", e)

def main():
    in_path = "decoder_in"
    out_path = "decoder_out"
    if not os.path.isdir(in_path):
        os.makedirs(in_path)
        print("Put all files to extract in decoder_in folder")
        input("Press Enter/Return to exit")
        return
        
    os.makedirs(out_path, exist_ok=True)

    for _, _, files in os.walk(in_path):
        for name in files:
            print(name)
            convert(in_path, out_path, name)

main()