import torch
from collections import defaultdict, deque
from typing import Dict, Any, List
from comfy_integration.nodes import NODE_CLASS_MAPPINGS
from utils.app_utils import get_value_at_index


def _normalize_dynamic_inputs(node_class, class_type, kwargs):
    """Rebuild the structures ComfyUI's executor would have built for this node.

    Two V3 dynamic shapes need post-processing when nodes are called directly:

    * ``COMFY_DYNAMICCOMBO_V3`` — the widget arrives as a bare string but the
      node signature expects ``{id: <option>, **sub_keys}`` (e.g.
      ``sampling_mode={"sampling_mode": "off"}``).
    * ``COMFY_AUTOGROW_V3`` — slots arrive flattened as ``image1``, ``image2``,
      ... but the node signature expects ``images={"image1": <value>, ...}``.

    ComfyUI's ``comfy_api.latest._io.build_nested_inputs`` does both; this
    Space executes nodes directly and therefore has to redo the work.
    """
    schema = getattr(node_class, 'SCHEMA', None)
    inputs = getattr(schema, 'inputs', None) if schema is not None else None
    if not inputs:
        return kwargs

    # --- COMFY_AUTOGROW_V3 first: consume image1/video1/... slots ---
    for spec in inputs:
        sub_id = getattr(spec, 'id', None)
        template = getattr(spec, 'template', None)
        if not sub_id or template is None:
            continue

        names = getattr(template, 'names', None) or []
        prefix = getattr(template, 'prefix', None)
        candidates = list(names)
        if prefix:
            max_n = getattr(template, 'max', 0) or 0
            if not candidates:
                candidates = [f"{prefix}{i}" for i in range(max_n)]

        collected = {}
        for name in candidates:
            if name in kwargs:
                collected[name] = kwargs.pop(name)

        if collected and not isinstance(kwargs.get(sub_id), dict):
            kwargs[sub_id] = collected

    # --- COMFY_DYNAMICCOMBO_V3 ---
    for spec in inputs:
        sub_id = getattr(spec, 'id', None)
        options = getattr(spec, 'options', None)
        if not sub_id or not options or sub_id not in kwargs:
            continue

        value = kwargs[sub_id]
        if isinstance(value, dict):
            continue  # already structured

        option_keys = []
        for opt in options:
            if isinstance(opt, dict):
                option_keys.append(opt.get('key'))
            else:
                option_keys.append(getattr(opt, 'key', None))

        if value not in option_keys:
            print(f"Warning: '{sub_id}' value {value!r} is not one of {option_keys} for {class_type}.")
            continue

        # Pull sibling sub-keys (e.g. "sampling_mode.temperature") into the dict.
        nested = {sub_id: str(value)}
        prefix = f"{sub_id}."
        for key in list(kwargs.keys()):
            if key.startswith(prefix):
                nested[key[len(prefix):]] = kwargs.pop(key)
        kwargs[sub_id] = nested
        print(f"  -> Normalized dynamic input '{sub_id}' for {class_type}: {nested}")

    return kwargs


class WorkflowExecutor:
    @staticmethod
    def topological_sort(workflow: Dict[str, Any]) -> List[str]:
        graph = defaultdict(list)
        in_degree = {node_id: 0 for node_id in workflow}

        for node_id, node_info in workflow.items():
            for input_value in node_info.get('inputs', {}).values():
                if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
                    source_node_id = input_value[0]
                    if source_node_id in workflow:
                        graph[source_node_id].append(node_id)
                        in_degree[node_id] += 1

        queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0])
        
        sorted_nodes = []
        while queue:
            current_node_id = queue.popleft()
            sorted_nodes.append(current_node_id)

            for neighbor_node_id in graph[current_node_id]:
                in_degree[neighbor_node_id] -= 1
                if in_degree[neighbor_node_id] == 0:
                    queue.append(neighbor_node_id)
        
        if len(sorted_nodes) != len(workflow):
            raise RuntimeError("Workflow contains a cycle and cannot be executed.")
            
        return sorted_nodes

    @staticmethod
    def execute_workflow(workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
        with torch.no_grad():
            computed_outputs = initial_objects
            
            try:
                sorted_node_ids = WorkflowExecutor.topological_sort(workflow)
                
                final_node_id = None
                for node_id in reversed(sorted_node_ids):
                    if workflow[node_id].get('class_type') == 'SaveImage':
                        final_node_id = node_id
                        break
                
                if final_node_id:
                    required_nodes = set()
                    nodes_to_visit = [final_node_id]
                    while nodes_to_visit:
                        curr_id = nodes_to_visit.pop()
                        if curr_id in required_nodes:
                            continue
                        required_nodes.add(curr_id)
                        curr_info = workflow.get(curr_id, {})
                        for input_val in curr_info.get('inputs', {}).values():
                            if isinstance(input_val, list) and len(input_val) == 2 and isinstance(input_val[0], str):
                                src_id = input_val[0]
                                if src_id in workflow and src_id not in required_nodes:
                                    nodes_to_visit.append(src_id)
                    
                    sorted_node_ids = [nid for nid in sorted_node_ids if nid in required_nodes]

                print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}")
            except RuntimeError as e:
                print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---")
                for node_id, node_info in workflow.items():
                    print(f"  Node {node_id} ({node_info['class_type']}):")
                    for input_name, input_value in node_info['inputs'].items():
                         if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
                             print(f"    - {input_name} <- [{input_value[0]}, {input_value[1]}]")
                raise e

            for node_id in sorted_node_ids:
                if node_id in computed_outputs:
                    continue
                    
                node_info = workflow[node_id]
                class_type = node_info['class_type']
                
                is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs'])
                if node_id in initial_objects and is_loader_with_filename:
                    continue
                
                node_class = NODE_CLASS_MAPPINGS.get(class_type)
                if node_class is None:
                     raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?")
                
                node_instance = node_class()
                
                kwargs = {}
                for param_name, param_value in node_info['inputs'].items():
                    if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str):
                        source_node_id, output_index = param_value
                        if source_node_id not in computed_outputs:
                            raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.")
                        
                        source_output_tuple = computed_outputs[source_node_id]
                        actual_value = get_value_at_index(source_output_tuple, output_index)
                    else:
                        actual_value = param_value

                    if '.' in param_name:
                        parent_key, child_key = param_name.split('.', 1)
                        if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict):
                            kwargs[parent_key] = {}
                        kwargs[parent_key][child_key] = actual_value
                    else:
                        kwargs[param_name] = actual_value

                # Safety guard for nodes expecting dynamic autogrow dictionary (e.g. TextEncodeMingImageEdit)
                if class_type == 'TextEncodeMingImageEdit' and ('images' not in kwargs or kwargs.get('images') is None):
                    kwargs['images'] = {}

                # ComfyUI's own executor restructures V3 dynamic inputs into the
                # nested dict the node signature expects (see
                # comfy_api.latest._io.build_nested_inputs). This direct-call
                # path bypasses that step, so a COMFY_DYNAMICCOMBO_V3 widget
                # arrives as a bare string and the node fails on
                # `sampling_mode.get(...)`. Rebuild it here.
                kwargs = _normalize_dynamic_inputs(node_class, class_type, kwargs)

                function_name = getattr(node_class, 'FUNCTION')
                execution_method = getattr(node_instance, function_name)

                result = execution_method(**kwargs)
                computed_outputs[node_id] = result
            
            final_node_id = None
            for node_id in reversed(sorted_node_ids):
                 if workflow[node_id]['class_type'] == 'SaveImage':
                     final_node_id = node_id
                     break
            
            if not final_node_id:
                raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.")

            save_image_inputs = workflow[final_node_id]['inputs']
            image_source_node_id, image_source_index = save_image_inputs['images']
            
            return get_value_at_index(computed_outputs[image_source_node_id], image_source_index)