import os
import sys
from PIL import Image  # type: ignore
import time
import concurrent.futures
from typing import Dict, List, Optional, Tuple

# Define suffix lists for BaseColor, Normal, packed maps, and single-channel maps
BASECOLOR_SUFFIXES = ['_alb.', '_albedo.', '_bc.', '_basecolor.', '_b.']
NORMAL_SUFFIXES = ['_nrm.', '_normal.', '_n.']
RMA_SUFFIXES = ['_rma.']
ORM_SUFFIXES = ['_orm.']
ROUGHNESS_SUFFIXES = ['_roughness.', '_rough.', '_rgh.']
METALLIC_SUFFIXES = ['_metallic.', '_metalness.', '_metal.', '_met.']
AO_SUFFIXES = ['_ao.', '_ambientocclusion.', '_occlusion.']
EMISSIVE_SUFFIXES = ['_emissive.']
OPACITY_SUFFIXES = ['_opacity.']
MASK_SUFFIXES = ['_mask.', '_m.']

def detect_texture_type(filename):
    """Detect the type of texture based on naming suffixes."""
    lowered = filename.lower()
    if any(suffix in lowered for suffix in BASECOLOR_SUFFIXES):
        return 'BaseColor'
    if any(suffix in lowered for suffix in NORMAL_SUFFIXES):
        return 'Normal'
    if any(suffix in lowered for suffix in RMA_SUFFIXES):
        return 'RMA'
    if any(suffix in lowered for suffix in ORM_SUFFIXES):
        return 'ORM'
    if any(suffix in lowered for suffix in ROUGHNESS_SUFFIXES):
        return 'Roughness'
    if any(suffix in lowered for suffix in METALLIC_SUFFIXES):
        return 'Metallic'
    if any(suffix in lowered for suffix in AO_SUFFIXES):
        return 'AO'
    if any(suffix in lowered for suffix in EMISSIVE_SUFFIXES):
        return 'Emissive'
    if any(suffix in lowered for suffix in OPACITY_SUFFIXES):
        return 'Opacity'
    if any(suffix in lowered for suffix in MASK_SUFFIXES):
        return 'Mask'
    return None

def get_material_name(filename):
    """Strip the T_/TX_ prefix while keeping the suffix for detection and return the material name."""
    base_name = os.path.basename(filename)

    if base_name.startswith('T_'):
        base_name = base_name[2:]
    elif base_name.startswith('TX_'):
        base_name = base_name[3:]

    return base_name.rsplit('_', 1)[0]

def convert_single_material(material_data: Tuple[str, Dict[str, str]], output_folder: str) -> Tuple[bool, str]:
    """Convert a single material to BCR/NMO format."""
    material, files = material_data
    basecolor_file = files.get('BaseColor')
    normal_file = files.get('Normal')
    rma_file = files.get('RMA')
    orm_file = files.get('ORM')
    roughness_file = files.get('Roughness')
    metallic_file = files.get('Metallic')
    ao_file = files.get('AO')
    emissive_file = files.get('Emissive')
    opacity_file = files.get('Opacity')
    mask_file = files.get('Mask')

    try:
        success, warnings = convert_to_bcr_nmo(
            material,
            basecolor_file,
            normal_file,
            rma_file,
            orm_file,
            roughness_file,
            metallic_file,
            ao_file,
            emissive_file,
            opacity_file,
            mask_file,
            output_folder,
        )
        if success:
            if warnings:
                return True, f"{material}: Successfully converted. Warning(s): {' '.join(warnings)}"
            return True, f"{material}: Successfully converted."
        return False, f"Skipping {material}: input file sizes do not match."
    except Exception as e:
        return False, f"Error processing {material}: {str(e)}"

def process_textures(input_files):
    """Main function to process all textures in a folder and convert to BCR/NMO."""
    textures: Dict[str, Dict[str, str]] = {}

    for filepath in input_files:
        filename = os.path.basename(filepath)
        material_name = get_material_name(filename)
        texture_type = detect_texture_type(filename)

        if texture_type is None:
            continue

        if material_name not in textures:
            textures[material_name] = {}
        textures[material_name][texture_type] = filepath

    base_path = os.path.dirname(input_files[0])
    output_folder = os.path.join(base_path, 'merged')
    os.makedirs(output_folder, exist_ok=True)

    material_count = len(textures)
    print(f"Detected {material_count} Materials to process.")

    valid_materials: Dict[str, Dict[str, str]] = {}
    failed_converts = 0

    for material, files in textures.items():
        missing_files = []
        if not files.get('BaseColor'):
            missing_files.append('BaseColor')
        if not files.get('Normal'):
            missing_files.append('Normal')
        if not (files.get('RMA') or files.get('ORM')):
            if not files.get('Roughness'):
                missing_files.append('Roughness')
            if not files.get('AO'):
                missing_files.append('AO')

        if missing_files:
            print(f"Skipping {material}: missing {', '.join(missing_files)}")
            failed_converts += 1
        else:
            valid_materials[material] = files

    success_count = 0
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_material = {
            executor.submit(convert_single_material, (material, files), output_folder): material
            for material, files in valid_materials.items()
        }

        for future in concurrent.futures.as_completed(future_to_material):
            material = future_to_material[future]
            try:
                success, message = future.result()
                if success:
                    success_count += 1
                else:
                    failed_converts += 1
                print(message)
            except Exception as e:
                failed_converts += 1
                print(f"Error processing {material}: {str(e)}")

    print(f"+++{success_count} of {material_count} materials successfully converted+++")
    time.sleep(3)

def _ensure_single_channel(image_path: str):
    """Load an arbitrary image and return a single-channel version for merging."""
    channel_image = Image.open(image_path)
    if channel_image.mode != 'L':
        channel_image = channel_image.convert('L')
    return channel_image

def convert_to_bcr_nmo(
    material: str,
    basecolor_file: Optional[str],
    normal_file: Optional[str],
    rma_file: Optional[str],
    orm_file: Optional[str],
    roughness_file: Optional[str],
    metallic_file: Optional[str],
    ao_file: Optional[str],
    emissive_file: Optional[str],
    opacity_file: Optional[str],
    mask_file: Optional[str],
    output_folder: str,
) -> Tuple[bool, List[str]]:
    """Convert the provided textures to BCR and NMO targets."""
    if basecolor_file is None or normal_file is None:
        raise ValueError('BaseColor and Normal textures are required.')

    warnings: List[str] = []
    basecolor_img = Image.open(basecolor_file).convert('RGBA')
    normal_img = Image.open(normal_file).convert('RGBA')
    base_r, base_g, base_b, _ = basecolor_img.split()
    normal_r, normal_g, _, _ = normal_img.split()

    roughness_channel = metallic_channel = ao_channel = None

    if rma_file:
        packed_img = Image.open(rma_file).convert('RGBA')
        if not (basecolor_img.size == normal_img.size == packed_img.size):
            return False, warnings
        roughness_channel, metallic_channel, ao_channel, _ = packed_img.split()
    elif orm_file:
        packed_img = Image.open(orm_file).convert('RGBA')
        if not (basecolor_img.size == normal_img.size == packed_img.size):
            return False, warnings
        ao_channel, roughness_channel, metallic_channel, _ = packed_img.split()
    else:
        if roughness_file is None or ao_file is None:
            raise ValueError('Roughness and AO textures are required when RMA/ORM is absent.')
        roughness_channel = _ensure_single_channel(roughness_file)
        ao_channel = _ensure_single_channel(ao_file)
        if metallic_file is not None:
            metallic_channel = _ensure_single_channel(metallic_file)
        else:
            metallic_channel = Image.new('L', basecolor_img.size, 0)
            warnings.append(f"{material}: Missing Metallic texture. Using a black channel.")

        if not (
            basecolor_img.size
            == normal_img.size
            == roughness_channel.size
            == ao_channel.size
        ):
            return False, warnings
        if metallic_channel.size != basecolor_img.size:
            return False, warnings

    if roughness_channel is None or metallic_channel is None or ao_channel is None:
        raise RuntimeError('Failed to resolve packed or single-channel textures.')

    bcr_img = Image.merge('RGBA', (base_r, base_g, base_b, roughness_channel))
    bcr_img.save(os.path.join(output_folder, f"{material}_BCR.tga"))

    nmo_img = Image.merge('RGBA', (normal_r, normal_g, metallic_channel, ao_channel))
    nmo_img.save(os.path.join(output_folder, f"{material}_NMO.tga"))

    if emissive_file:
        emissive_img = Image.open(emissive_file)
        if emissive_img.mode != 'RGBA':
            emissive_img = emissive_img.convert('RGBA')
        emissive_img.save(os.path.join(output_folder, f"{material}_EM.tga"))

    if opacity_file:
        opacity_img = Image.open(opacity_file)
        if opacity_img.mode != 'RGBA':
            opacity_img = opacity_img.convert('RGBA')
        opacity_img.save(os.path.join(output_folder, f"{material}_OP.tga"))

    if mask_file:
        mask_img = Image.open(mask_file)
        if mask_img.mode != 'RGBA':
            mask_img = mask_img.convert('RGBA')
        mask_img.save(os.path.join(output_folder, f"{material}_MASK.tga"))

    return True, warnings

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: drag and drop texture files onto the script")
    else:
        input_files = sys.argv[1:]
        process_textures(input_files)