"""
    Utilities for Materials in Blender.
"""
import logging
import os
import copy
import random
from pathlib import Path
from typing import Tuple, Union

import bpy
import numpy as np

import gin
import zpy

log = logging.getLogger(__name__)


def verify(
    mat: Union[bpy.types.Material, str],
    check_none: bool = True,
) -> bpy.types.Material:
    """ Get a material given either its name or the object itself.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
        check_none (bool, optional): Check to make sure material exists. Defaults to True.

    Raises:
        ValueError: Material does not exist.

    Returns:
        bpy.types.Material: Material object.
    """
    if isinstance(mat, str):
        mat = bpy.data.materials.get(mat)
    if check_none and mat is None:
        raise ValueError(f'Could not find material {mat}.')
    return mat


def for_mat_in_obj(
    obj: Union[bpy.types.Object, str],
) -> bpy.types.Material:
    """ Yield materials in scene object.

    Args:
        obj (Union[bpy.types.Object, str]): Scene object (or it's name)

    Raises:
        ValueError: [description]

    Returns:
        bpy.types.Material: Material object.
    """
    obj = zpy.objects.verify(obj)
    if len(obj.material_slots) > 1:
        for mat in obj.material_slots:
            yield mat.material
    else:
        if obj.active_material is not None:
            return obj.active_material
        else:
            log.debug(
                f'No active material or material slots found for {obj.name}')
            return None


_SAVED_MATERIALS = {}


def save_mat_props(
    mat: Union[bpy.types.Material, str],
) -> None:
    """ Save a pose (rot and pos) to dict.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
    """
    log.info(f'Saving material properties for {mat.name}')
    _SAVED_MATERIALS[mat.name] = get_mat_props(mat)


def restore_mat_props(
    mat: Union[bpy.types.Material, str],
) -> None:
    """ Restore an object to a position.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
    """
    log.info(f'Restoring material properties for {mat.name}')
    set_mat_props(mat, _SAVED_MATERIALS[mat.name])


def restore_all_mat_props() -> None:
    """ Restore all jittered materials to original look. """
    for mat_name, mat_props in _SAVED_MATERIALS.items():
        set_mat_props(mat_name, mat_props)


def get_mat_props(
    mat: Union[bpy.types.Material, str],
) -> Tuple[float]:
    """ Get (some of the) material properties.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)

    Returns:
        Tuple[float]: Material property values (roughness, metallic, specular).
    """
    mat = verify(mat)
    bsdf_node = mat.node_tree.nodes.get('Principled BSDF')
    if bsdf_node is None:
        log.warning(f'No BSDF node in {mat.name}')
        return (0.0, 0.0, 0.0)
    return (
        bsdf_node.inputs['Roughness'].default_value,
        bsdf_node.inputs['Metallic'].default_value,
        bsdf_node.inputs['Specular'].default_value,
    )


def set_mat_props(
    mat: Union[bpy.types.Material, str],
    prop_tuple: Tuple[float],
) -> None:
    """ Set (some of the) material properties.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
        prop_tuple (Tuple[float]): Material property values (roughness, metallic, specular).
    """
    mat = verify(mat)
    # TODO: Work backwards from Material output node instead of
    #       assuming a 'Principled BSDF' node
    bsdf_node = mat.node_tree.nodes.get('Principled BSDF', None)
    if bsdf_node is None:
        log.warning(f'No BSDF node in {mat.name}')
        return
    bsdf_node.inputs['Roughness'].default_value = copy.copy(prop_tuple[0])
    bsdf_node.inputs['Metallic'].default_value = copy.copy(prop_tuple[1])
    bsdf_node.inputs['Specular'].default_value = copy.copy(prop_tuple[2])


@gin.configurable
def jitter(
    mat: Union[bpy.types.Material, str],
    std: float = 0.2,
    save_first_time: bool = True,
) -> None:
    """ Randomize an existing material a little.

    Args:
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
        std (float, optional): Standard deviation of gaussian noise over material property. Defaults to 0.2.
        save_first_time (bool, optional): Save the material props first time jitter is called and restore before jittering every subsequent time. Defaults to True.
    """
    mat = verify(mat)
    if save_first_time:
        if _SAVED_MATERIALS.get(mat.name, None) is None:
            save_mat_props(mat)
        else:
            restore_mat_props(mat)
    log.info(f'Jittering material {mat.name}')
    mat_props = get_mat_props(mat)
    jittered_mat_props = tuple(
        map(lambda p: p + random.gauss(0, std), mat_props))
    set_mat_props(mat, jittered_mat_props)


@gin.configurable
def random_texture_mat(
    texture_dir: Union[Path, str] = 'lib/textures/random_512p',
    relative_to_assets_dir: bool = True,
) -> bpy.types.Material:
    """ Generate a random material from a directory of random texture images.

    Args:
        texture_dir (Union[Path, str], optional): Path to directory with texture images.
        relative_to_assets_dir (bool, optional): Path is relative to the $ASSETS directory. Defaults to False.

    Returns:
        bpy.types.Material: The newly created material.
    """
    if relative_to_assets_dir:
        texture_dir = zpy.blender.get_asset_lib_path().joinpath(texture_dir)
    texture_dir = zpy.files.verify_path(texture_dir, check_dir=True)
    # Create list of texture images in directory
    texture_paths = []
    for _path in texture_dir.iterdir():
        if _path.is_file() and _path.suffix in ['.jpg', '.png']:
            texture_paths.append(_path)
    texture_path = random.choice(texture_paths)
    log.info(f'Found {len(texture_paths)} Textures at {texture_dir}')
    log.info(f'Randomly picked {texture_path.stem}')
    return make_mat_from_texture(texture_path, name=texture_path.stem)


@gin.configurable
def make_mat_from_texture(
    texture_path: Union[Path, str],
    name: str = None,
) -> bpy.types.Material:
    """ Makes a material from a texture image.

    Args:
        texture_path (Union[Path, str]): Path to texture image.
        name (str, optional): Name of new material.

    Returns:
        bpy.types.Material: The newly created material.
    """
    texture_path = zpy.files.verify_path(texture_path, make=False)
    if name is None:
        name = texture_path.stem
    mat = bpy.data.materials.get(name, None)
    if mat is None:
        log.debug(f'Material {name} does not exist, creating it.')
        mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    bsdf_node = mat.node_tree.nodes.get('Principled BSDF')
    out_node = mat.node_tree.nodes.get('Material Output')
    tex_node = mat.node_tree.nodes.new('ShaderNodeTexImage')
    tex_node.name = 'ImageTexture'
    coord_node = mat.node_tree.nodes.new('ShaderNodeTexCoord')
    bpy.ops.image.open(filepath=str(texture_path))
    tex_node.image = bpy.data.images[texture_path.name]
    tex_node.image.colorspace_settings.name = 'Filmic Log'
    mat.node_tree.links.new(tex_node.outputs[0], bsdf_node.inputs[0])
    mat.node_tree.links.new(coord_node.outputs[2], tex_node.inputs[0])
    mat.node_tree.links.new(out_node.inputs[0], bsdf_node.outputs[0])
    tex_node.image.reload()
    return mat


@gin.configurable
def make_mat_from_color(
    color: Tuple[float],
    name: str = None,
) -> bpy.types.Material:
    """ Makes a material given a color.

    Args:
        color (Tuple[float]): Color tuple (RGB).
        name (str, optional): Name of new material.

    Returns:
        bpy.types.Material: The newly created material.
    """
    if name is None:
        name = str(color)
    mat = bpy.data.materials.get(name, None)
    if mat is None:
        log.debug(f'Material {name} does not exist, creating it.')
        mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    bsdf_node = mat.node_tree.nodes.get('Principled BSDF')
    out_node = mat.node_tree.nodes.get('Material Output')
    mat.node_tree.nodes.remove(bsdf_node)
    bsdf_node = mat.node_tree.nodes.new('ShaderNodeBsdfDiffuse')
    bsdf_node.inputs['Color'].default_value = color + (1.,)
    mat.node_tree.links.new(out_node.inputs[0], bsdf_node.outputs[0])
    return mat


def set_mat(
    obj: Union[bpy.types.Object, str],
    mat: Union[bpy.types.Material, str],
    recursive: bool = True,
) -> None:
    """ Set the material for an object.

    Args:
        obj (Union[bpy.types.Object, str]): Scene object (or it's name) with an active material.
        mat (Union[bpy.types.Material, str]):  Material (or it's name)
        recursive (bool, optional): Recursively set material for child objects. Defaults to True.
    """
    obj = zpy.objects.verify(obj)
    mat = zpy.material.verify(mat)
    if hasattr(obj, 'active_material'):
        log.debug(f'Setting object {obj.name} material {mat.name}')
        obj.active_material = mat
    else:
        log.warning('Object does not have material property')
        return
    # Recursively change material on all children of object
    if recursive:
        for child in obj.children:
            set_mat(child, mat)


@gin.configurable
def make_aov_material_output_node(
    mat: bpy.types.Material = None,
    obj: bpy.types.Object = None,
    style: str = 'instance',
) -> None:
    """ Make AOV Output nodes in Composition Graph.

    Args:
        mat (bpy.types.Material, optional): A blender material (either it's name or the object itself).
        obj (bpy.types.Object, optional): A blender object (either it's name or the object itself).
        style (str, optional): Type of segmentation in [instance, category]. Defaults to 'instance'.

    Raises:
        ValueError: Invalid style, no object or material given.
    """
    # Make sure engine is set to Cycles
    scene = zpy.blender.verify_blender_scene()
    if not (scene.render.engine == "CYCLES"):
        log.warning(' Setting render engine to CYCLES to use AOV')
        scene.render.engine == "CYCLES"

    # TODO: Refactor this legacy "styles" code

    # Only certain styles are available
    valid_styles = ['instance', 'category']
    assert style in valid_styles, \
        f'Invalid style {style} for AOV material output node, must be in {valid_styles}.'

    # HACK: multiple material slots
    all_mats = []

    # Use material
    if mat is not None:
        all_mats = [mat]
    # Get material from object
    elif obj is not None:
        if obj.active_material is None:
            log.debug(f'No active material found for {obj.name}')
            return
        if len(obj.material_slots) > 1:
            for mat in obj.material_slots:
                all_mats.append(mat.material)
        else:
            all_mats.append(obj.active_material)
    else:
        raise ValueError('Must pass in an Object or Material')

    # HACK: multiple material slots
    for mat in all_mats:

        # Make sure material is using nodes
        if not mat.use_nodes:
            mat.use_nodes = True
        tree = mat.node_tree

        # Vertex Color Node
        vcol_node = zpy.nodes.get_or_make(
            f'{style} Vertex Color', 'ShaderNodeVertexColor', tree)
        vcol_node.layer_name = style

        # AOV Output Node
        # HACK: This type of node has a "name" property which prevents using the
        # normal zpy.nodes code due to a scope conflict with the bpy.types.Node.name property
        # See: https://docs.blender.org/api/current/bpy.types.ShaderNodeOutputAOV.html
        _name = style
        aovout_node = None
        for _node in tree.nodes:
            if _node.name == _name:
                aovout_node = _node
        if aovout_node is None:
            aovout_node = tree.nodes.new('ShaderNodeOutputAOV')
        aovout_node.name = style

        tree.links.new(vcol_node.outputs['Color'], aovout_node.inputs['Color'])
