from PIL import Image
import sys
import os
import time
import numpy as np
from multiprocessing import Pool, cpu_count, freeze_support
from concurrent.futures import ThreadPoolExecutor
import threading

def process_image(args):
    input_path, export_dir, idx, total = args
    try:
        # Open image and check if it has alpha
        currentTex = Image.open(input_path)
        has_alpha = currentTex.mode == 'RGBA'
        was_resized = False
        
        if has_alpha:
            # Get the alpha channel
            r, g, b, a = currentTex.split()
            # Convert alpha to numpy array to check its values
            alpha_array = np.array(a)
            # Check if alpha channel is useful (contains values other than 0 or 255)
            has_useful_alpha = np.any((alpha_array > 0) & (alpha_array < 255))
            
            if has_useful_alpha:
                # If alpha is useful, keep it
                if currentTex.size > (2048, 2048):
                    resizedTex = currentTex.resize((2048, 2048), Image.LANCZOS)
                    was_resized = True
                else:
                    resizedTex = currentTex
            else:
                # If alpha is not useful (all black or all white), ignore it
                currentTex = currentTex.convert('RGB')
                if currentTex.size > (2048, 2048):
                    resizedTex = currentTex.resize((2048, 2048), Image.LANCZOS)
                    was_resized = True
                else:
                    resizedTex = currentTex
                # Create new RGBA with full opacity
                rgba_image = Image.new('RGBA', resizedTex.size)
                rgba_image.paste(resizedTex, (0, 0))
                resizedTex = rgba_image
        else:
            # No alpha channel, just resize if needed
            if currentTex.size > (2048, 2048):
                resizedTex = currentTex.resize((2048, 2048), Image.LANCZOS)
                was_resized = True
            else:
                resizedTex = currentTex
            # Convert to RGBA with full opacity
            rgba_image = Image.new('RGBA', resizedTex.size)
            rgba_image.paste(resizedTex, (0, 0))
            resizedTex = rgba_image
        
        # Change the file extension to .tga
        output_filename = os.path.splitext(os.path.basename(input_path))[0] + '.tga'
        output_path = os.path.join(export_dir, output_filename)
        
        # Save as TGA format with RLE compression
        resizedTex.save(output_path, format='TGA', rle=True)
        
        # Print appropriate message based on whether the file was resized
        if was_resized:
            print(f"{os.path.basename(input_path)} optimized and converted to TGA ({idx}/{total})")
        else:
            print(f"{os.path.basename(input_path)} converted to TGA ({idx}/{total})")
    except Exception as e:
        print(f"Error processing {input_path} ({idx}/{total}): {e}")

def main():
    if len(sys.argv) < 2:
        print("Please provide the input files.")
        sys.exit(1)

    inputTexs = sys.argv[1:]
    export_dir = os.path.join(os.path.dirname(inputTexs[0]), "optimized")

    if not os.path.exists(export_dir):
        os.mkdir(export_dir)

    total_files = len(inputTexs)
    print(f"Starting optimization of {total_files} files")

    # Start timing
    start_time = time.perf_counter()

    args_list = [
        (input_path, export_dir, idx + 1, total_files)
        for idx, input_path in enumerate(inputTexs)
    ]

    # Check if running as a frozen executable
    if getattr(sys, 'frozen', False):
        # Use ThreadPoolExecutor when frozen
        max_workers = min(32, (os.cpu_count() or 1) + 4)
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            list(executor.map(process_image, args_list))
    else:
        # Use multiprocessing when running from console
        with Pool(processes=cpu_count()) as pool:
            pool.map(process_image, args_list)

    # End timing
    end_time = time.perf_counter()
    elapsed_time = end_time - start_time

    print(f"Textures successfully resized in {elapsed_time:.2f} seconds!")
    time.sleep(5)

if __name__ == "__main__":
    freeze_support()
    main()