# Copyright (C) 2022 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.

import os

from typing import List, Tuple, Dict

from .pbr_entry import PBREntry


_img_ext = set('.jpg,.jpeg,.png,.tga,.tif,.tiff,.gif,.hdr,.exr'.split(','))

_ignore = set('xtm'.split(','))
_variant_suffixes = set('05k,1k,2k,4k,8k'.split(','))
_replaces = [('nor_dx', 'nor-dx'), ('nor_gl', 'nor-gl')]

_diffuse_ext = set('basecolor,base_color,diffuse,diff,albedo,color,col'.split(','))
_spec_ext = set('specularlevel,specular,spec,spc'.split(','))
_rough_ext = set('roughness,rough,rgh'.split(','))
_gloss_ext = set('gloss,gls'.split(','))
_normal_ext = set('nor-dx,nor-gl,nor_gl,normal,norm,nor,nrm'.split(','))
_metal_ext = set('metallic,metal,met,metalness'.split(','))
_height_ext = set('height,hgt,bump'.split(','))
_disp_ext = set('displacement,displace,disp,dis'.split(','))
_emission_ext = set('emit,emission,emissive'.split(','))
_alpha_ext = set('opacity,alpha,mask'.split(','))
_ao_ext = set('ambientocclusion,ao,colao'.split(','))
_arm_ext = set('arm'.split(','))


def is_image(extension: str):
    """
    Check if the extension is a supported image.
    """
    return extension in _img_ext


def scan_sets(files: List[str]) -> List[PBREntry]:
    entries = [] # type: List[Tuple[str, List[str]]]

    # Create format that is better for analyzing.
    for f in files:
        # Extract filename and extension.
        base, ext = os.path.splitext(os.path.split(f)[1])

        # Only handle images ..
        if ext in _img_ext:
            # Some use '-' as delimiter, bring to one format.
            base = base.lower().replace('-', '_')
            # Some channel specifiers contain '_', like nor_gl,
            # transform _ to - for further processing.
            for r in _replaces:
                base = base.replace(*r)
            # Add full image file path and splitted name.
            entries.append((f, reversed(base.split('_'))))


    # Collect all sets temporary in this dict.
    temp_collector = {} # type: Dict[str, Dict[str, str]]
    for f, parts in entries:
        # If no variant tag is given, use 'Default'
        variant = 'Default'

        # Collect all name parts before channel specifier in this list.
        new_parts = []

        # Eventually store channel for this image ..
        channel = None
        # Loop over name parts (they are reversed here (see insertion)).
        for p in parts:
            # Handle ignore and variant first.
            if p in _ignore: 
                pass
            else:
                # Channel is extracted only once (so the last one).
                if not channel:
                    if p in _diffuse_ext: channel = 'diffuse'
                    elif p in _spec_ext: channel = 'specular'
                    elif p in _rough_ext: channel = 'roughness'
                    elif p in _gloss_ext: channel = 'gloss'
                    elif p in _normal_ext: channel = 'normal'
                    elif p in _metal_ext: channel = 'metal'
                    elif p in _height_ext: channel = 'height'
                    elif p in _disp_ext: channel = 'displace'
                    elif p in _emission_ext: channel = 'emission'
                    elif p in _alpha_ext: channel = 'alpha'
                    elif p in _ao_ext: channel = 'ao'
                    elif p in _arm_ext: channel = 'arm'
                    elif p == 'preview': channel = 'preview'
                    elif p in _variant_suffixes: variant = p
                else:
                    # Recollect set name, only if channel is known, we scan from back to front.
                    if channel:
                        new_parts.append(p)

        # Known channel must be found, othwise this image isn't supported.
        if channel and channel != 'preview':
            # Add this in one set inside the temp collector.
            pbr = '-'.join(reversed(new_parts)) + f'_{variant}'
            temp_collector.setdefault(pbr, dict())[channel] = f
        else:
            if channel != 'preview':
                pass #print(f'Unknown: {f}')

    # Now convert all entries to PBR sets.
    r = [] # type: List[PBREntry]
    for k, v in temp_collector.items():
        # Diffuse is a must have.
        if 'diffuse' in v:
            # We need key to extract name and variant.
            basename, variant = k.split('_')
            name = os.path.split(v['diffuse'])[1][:len(basename)]
            r.append(
                PBREntry(
                    name=name,
                    variant_name=variant,
                    diffuse=v.get('diffuse'),
                    alpha=v.get('alpha'),
                    ao=v.get('ao'),
                    specular=v.get('specular'),
                    metal=v.get('metal'),
                    roughness=v.get('roughness'),
                    glossiness=v.get('gloss'),
                    normal=v.get('normal'),
                    height=v.get('height'),
                    displace=v.get('displace'),
                    emission=v.get('emission'),
                    arm=v.get('arm'),
                )
            )
        
    return r

