import bpy
import os
import re
import json

from .constant import TECH_PASS_KEYWORDS

### --- Manage nodes --- ###

def real_loc(n):
    if not n.parent:
        return n.location
    return n.location + real_loc(n.parent)

def set_file_output_format(fo):
    """Default fileout format for output file node
    Get from OUTPUT_RENDER_FILE_FORMAT environment variable
    else fallback to multilayer EXR
    """

    env_file_format = json.loads(os.environ.get('OUTPUT_RENDER_FILE_FORMAT', '{}'))
    if not env_file_format:
        env_file_format = {
            'file_format': 'OPEN_EXR_MULTILAYER',
            'exr_codec': 'ZIP',
            'color_depth': '16',
            'color_mode': 'RGBA'
        }

    for k, v in env_file_format.items():
        setattr(fo.format, k, v)

def clear_disconnected(fo):
    '''Remove unlinked inputs from file output node'''
    for inp in reversed(fo.inputs):
        if not inp.is_linked:
            print(f'Deleting unlinked fileout slot: {inp.name}')
            fo.inputs.remove(inp)

def create_node(type, tree=None, **kargs):
    '''Get a type, a tree to add in, and optionnaly multiple attribute to set
    return created node
    '''
    tree = tree or bpy.context.scene.node_tree

    node = tree.nodes.new(type)
    for k,v in kargs.items():
        setattr(node, k, v)

    return node

def recursive_node_connect_check(l, target_node):
    '''Get a link and a node    
    return True if link is connected to passed target_node further in the node tree
    '''
    if l.to_node == target_node:
        return True    
    for o in l.to_node.outputs:
        for sl in o.links:
            if recursive_node_connect_check(sl, target_node):
                return True
    return False


def create_and_connect_file_output(node, outputs, file_out, out_name, base_path, out_base, 
                                   scene, remap_names, file_format, suffix='', color_depth=None, 
                                   location_offset=(500, 50)):
    """Helper function to create file output node and connect it to outputs
    
    Args:
        node: Source node
        outputs: List of outputs to connect
        file_out: Existing file output node or None
        out_name: Name for the file output node
        base_path: Base path for output files
        out_base: Base name for output files
        scene: Blender scene
        remap_names: Dictionary for remapping output names
        file_format: Format settings
        suffix: Optional suffix for paths
        color_depth: Optional override for color depth
        location_offset: Offset for the node position
        
    Returns:
        bpy.types.CompositorNode: Created or used file output node
    """
    if not outputs:
        return None
        
    nodes = scene.node_tree.nodes
    links = scene.node_tree.links
    
    fo = file_out
    if not fo:
        fo = nodes.get(out_name)
    if not fo:
        fo = create_node('CompositorNodeOutputFile', tree=scene.node_tree, 
                         location=(real_loc(node)[0] + location_offset[0], 
                                  real_loc(node)[1] + location_offset[1]), width=600)
        fo.inputs.remove(fo.inputs[0])  # Remove default image input
        
        if file_format:
            for k, v in file_format.items():
                setattr(fo.format, k, v)
        else:
            set_file_output_format(fo)
            
        if color_depth:
            fo.format.color_depth = color_depth
            
        fo.name = out_name
        if node.parent:
            fo.parent = node.parent

        if base_path:
            fo.base_path = base_path
        else:
            if fo.format.file_format == 'OPEN_EXR_MULTILAYER':
                fo.base_path = f'//render/{out_base}/{suffix}{out_base}_'
            else:
                fo.base_path = f'//render/{out_base}/{suffix}'

    for o in outputs:
        if next((l for l in o.links if recursive_node_connect_check(l, fo)), None):
            continue
        
        if (socket_remaps := remap_names.get(node.name)) and (custom_name := socket_remaps.get(o.name)):
            slot_name = bpy.path.clean_name(custom_name)
        else:
            slot_name = bpy.path.clean_name(o.name)

        fs = fo.file_slots.new('tmp')
        ls = fo.layer_slots.new('tmp')

        ls = fo.layer_slots[-1]
        ls.name = slot_name

        fs = fo.file_slots[-1]
        fs.path = f'{slot_name}/{slot_name}_'
        
        out_input = fo.inputs[-1]
        links.new(o, out_input)

    clear_disconnected(fo)
    fo.update()
    return fo


def connect_to_file_output(node_list, file_out=None, base_path='', excludes=None, remap_names=None, 
                           file_format=None, split_tech_passes=False):
    """Connect selected nodes output to file output(s)
    if a file output is selected, add intputs on it

    Args:
        node_list (list[bpy.types.Nodes,]): Nodes to connect
    
        file_out (bpy.types.CompositorNode, optional): File output node to connect to instead of new
            Defaults to None
        base_path (str, optional): Directory of images to render.
            if not passed, will use source node layer name
            Defaults to ''.

        file_format (dict, optionnal): converts each dictionary key into a file output format 
            attribute and assigns the corresponding value.
            Defaults to None.

        excludes (dict, optionnal): List of output names to exclude {node_name: [outputs,]}.
            Defaults toNone.

        remap_names (dict, optionnal): List of output names to remap {node_name: {output_name: new_name}}.
            
        split_tech_passes (bool, optional): When True, create a separate file output for technical passes
            Defaults to False.

    Returns:
        list[bpy.types.CompositorNode]: All nodes created.
    """

    scene = bpy.context.scene
    if not isinstance(node_list, list):
        node_list = [node_list]
    node_list = [n for n in node_list if n.type != 'OUTPUT_FILE']
    if not node_list:
        return []

    created_nodes = []
    excludes = excludes or {}
    remap_names = remap_names or {}


    tech_file_format = file_format.copy() if file_format else None

    if tech_file_format:
        ## TODO: Set to EXR if provioded is not ERX or EXR multilayer ?
        # ->> if tech_file_format.get('file_format')) is not in ('OPEN_EXR', 'OPEN_EXR_MULTILAYER'): ...

        ## Force 32 bit
        tech_file_format['color_depth'] = '32'
        ## Enforce a lossless format if provided is not
        if (codec := tech_file_format.get('exr_codec')) and codec not in ['ZIP','PIZ','RLE','ZIPS']:
            tech_file_format['exr_codec'] = 'PIZ'

    for node in node_list:
        exclusions = excludes.get(node.name) or []
        
        # Get all available outputs excluding those in exclusions
        all_outputs = [o for o in node.outputs if not o.is_unavailable and o.name not in exclusions]
        
        # Base name for output nodes
        if node.label:
            out_base = node.label
        else:
            out_base = node.name
        out_base = bpy.path.clean_name(out_base)
        
        # Categorize outputs
        crypto_outputs = [o for o in all_outputs if 'crypto' in o.name.lower()]
        
        if split_tech_passes:
            # Filter tech passes
            tech_outputs = [o for o in all_outputs if o.name.lower() in TECH_PASS_KEYWORDS] # any(keyword in o.name.lower() for keyword in TECH_PASS_KEYWORDS)]
            # Regular outputs (excluding crypto and tech passes)
            regular_outputs = [o for o in all_outputs if o not in crypto_outputs and o not in tech_outputs]
        else:
            # If not splitting tech passes, include them with regular outputs
            regular_outputs = [o for o in all_outputs if o not in crypto_outputs]
            tech_outputs = []

        y_offset = 50
        node_margin = 100
        # Create and connect regular outputs
        if regular_outputs:
            out_name = f'OUT_{out_base}'
            fo_regular = create_and_connect_file_output(
                node, regular_outputs, file_out, out_name, base_path,
                out_base, scene, remap_names, file_format,
                location_offset=(500, y_offset)
            )
            if fo_regular and fo_regular not in created_nodes:
                created_nodes.append(fo_regular)


        y_offset += -22 * len(regular_outputs) - node_margin
        # Create and connect tech outputs with 32-bit depth if split_tech_passes is True
        if tech_outputs:
            out_name = f'OUT_{out_base}_tech'

            fo_tech = create_and_connect_file_output(
                node, tech_outputs, None, out_name, base_path,
                out_base, scene, remap_names, tech_file_format,
                suffix='tech/', color_depth='32',
                location_offset=(500, y_offset)
            )

            if fo_tech and fo_tech not in created_nodes:
                created_nodes.append(fo_tech)
            
            y_offset -= node_margin

        y_offset += -22 * len(tech_outputs)
        # Create and connect crypto outputs with 32-bit depth
        if crypto_outputs:
            out_name = f'OUT_{out_base}_cryptos'
    
            fo_crypto = create_and_connect_file_output(
                node, crypto_outputs, None, out_name, base_path,
                out_base, scene, remap_names, tech_file_format,
                suffix='cryptos/', color_depth='32',
                location_offset=(500, y_offset)
            )

            if fo_crypto and fo_crypto not in created_nodes:
                created_nodes.append(fo_crypto)
        
    
    return created_nodes