# Copyright 2025 Pavlo Penenko
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

"""
Type mappings between Metashade and MaterialX.
"""

# Canonical map from Metashade dtype class names to MaterialX type strings
_METASHADE_TO_MTLX = {
    'Float': 'float',
    'Int': 'integer',
    'RgbF': 'color3',
    'RgbaF': 'color4',
    'Float2': 'vector2',
    'Float3': 'vector3',
    'Float4': 'vector4',
    'Float3x3': 'matrix33',
    'Float4x4': 'matrix44',
    'Bool': 'boolean',
}

# Derive the inverse map from the canonical forward map
_MTLX_TO_METASHADE = {v: k for k, v in _METASHADE_TO_MTLX.items()}

# Add MaterialX type aliases (not in forward map, but valid MaterialX types)
_MTLX_TO_METASHADE.update({
    'int': 'Int',        # Alias for 'integer'
    'string': 'Int',     # Enums map to integers in shader generation
    # Standard MaterialX Typedefs/Structs
    'BSDF': 'BSDF',
    'EDF': 'EDF',
    'VDF': 'VDF',
    'surfaceshader': 'SurfaceShader',
    'volumeshader': 'VolumeShader',
    'displacementshader': 'DisplacementShader',
    'lightshader': 'LightShader',
    'material': 'Material',
})

# Types that don't have Metashade equivalents (for documentation)
_UNSUPPORTED_MTLX_TYPES = frozenset({
    'filename',     # Texture references
    'integerarray', # Arrays need special handling
    'floatarray',
})


def metashade_to_mtlx(dtype_factory):
    """Map a Metashade dtype factory to a MaterialX type string."""
    if dtype_factory is None:
        return None
    
    dtype = dtype_factory._get_dtype()
    dtype_name = dtype.__name__
    
    if dtype_name not in _METASHADE_TO_MTLX:
        raise ValueError(
            f"No MaterialX type mapping for Metashade dtype '{dtype_name}'"
        )
    
    return _METASHADE_TO_MTLX[dtype_name]

def mtlx_to_metashade_dtype(mtlx_type: str, sh):
    """
    Get the Metashade dtype factory for a MaterialX type.
    
    Args:
        mtlx_type: MaterialX type string (e.g., 'float', 'vector3')
        sh: The Metashade generator instance
        
    Returns:
        Dtype factory (e.g., sh.Float, sh.Float3) or None if not mappable
    """
    metashade_name = _MTLX_TO_METASHADE.get(mtlx_type)
    if metashade_name is None:
        return None
    
    return getattr(sh, metashade_name, None)


def register_mtlx_closure_structs(sh):
    """Declare MaterialX closure-related structs on a generator.
    
    These structs are required by BSDF closure-type nodes and must be
    registered before acquiring or wrapping such nodes.
    
    The structs are declared with ``emit=False`` because their definitions
    are already emitted by MaterialX's standard library includes.
    
    Args:
        sh: The Metashade generator instance.
    """
    sh.struct('ClosureData', emit=False)(
        closureType=sh.Int,
        L=sh.Vector3f,
        V=sh.Vector3f,
        N=sh.Vector3f,
        P=sh.Point3f,
        occlusion=sh.Float
    )
    
    sh.struct('BSDF', emit=False)(
        response=sh.Float3,
        throughput=sh.Float3
    )

