"""
Control flow nodes for conditional logic and branching.
"""

from typing import Any, Dict, List
from ybagent.node import BaseNode, NodeType, NodeInput, NodeOutput


class ConditionalNode(BaseNode):
    """
    Conditional node for if/else branching.
    """
    
    def __init__(self, id: str, condition: str = "", **config):
        """
        Initialize conditional node.
        
        Args:
            id: Node ID
            condition: Python expression to evaluate (e.g., "value > 10")
            **config: Additional configuration
        """
        config["condition"] = condition
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.CONDITIONAL, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="input", type="any", required=True, description="Input data"),
            NodeInput(name="value", type="any", required=False, description="Value to test")
        ]
        self.metadata.outputs = [
            NodeOutput(name="true", type="any", description="Output if condition is true"),
            NodeOutput(name="false", type="any", description="Output if condition is false"),
            NodeOutput(name="result", type="bool", description="Condition result")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Evaluate condition and route output.
        
        Args:
            inputs: Input data and value to test
            context: Execution context
        
        Returns:
            Outputs based on condition result
        """
        condition = self.metadata.config.get("condition", "")
        if not condition:
            raise ValueError("Condition not specified")
        
        # Prepare evaluation context
        eval_context = {
            "input": inputs.get("input"),
            "value": inputs.get("value", inputs.get("input")),
        }
        
        # Evaluate condition
        try:
            result = eval(condition, {"__builtins__": {}}, eval_context)
        except Exception as e:
            raise ValueError(f"Error evaluating condition '{condition}': {e}")
        
        # Route output based on result
        output_data = inputs.get("input")
        
        return {
            "result": bool(result),
            "true": output_data if result else None,
            "false": output_data if not result else None,
            "output": output_data
        }


class SwitchNode(BaseNode):
    """
    Switch node for multi-way branching.
    """
    
    def __init__(self, id: str, cases: Dict[str, str] = None, **config):
        """
        Initialize switch node.
        
        Args:
            id: Node ID
            cases: Dictionary of case_name: condition
            **config: Additional configuration
        """
        config["cases"] = cases or {}
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.SWITCH, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="input", type="any", required=True, description="Input data"),
            NodeInput(name="value", type="any", required=False, description="Value to test")
        ]
        
        # Dynamic outputs based on cases
        cases = self.metadata.config.get("cases", {})
        outputs = [
            NodeOutput(name=case_name, type="any", description=f"Output for {case_name}")
            for case_name in cases.keys()
        ]
        outputs.append(NodeOutput(name="default", type="any", description="Default output"))
        self.metadata.outputs = outputs
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Evaluate cases and route to matching output.
        
        Args:
            inputs: Input data
            context: Execution context
        
        Returns:
            Outputs based on matching case
        """
        cases = self.metadata.config.get("cases", {})
        input_data = inputs.get("input")
        test_value = inputs.get("value", input_data)
        
        # Prepare evaluation context
        eval_context = {
            "input": input_data,
            "value": test_value,
        }
        
        # Evaluate each case
        matched_case = None
        for case_name, condition in cases.items():
            try:
                if eval(condition, {"__builtins__": {}}, eval_context):
                    matched_case = case_name
                    break
            except Exception as e:
                raise ValueError(f"Error evaluating case '{case_name}': {e}")
        
        # Build output
        result = {"default": None}
        for case_name in cases.keys():
            result[case_name] = None
        
        if matched_case:
            result[matched_case] = input_data
        else:
            result["default"] = input_data
        
        result["output"] = input_data
        return result


class LoopNode(BaseNode):
    """
    Loop node for iteration.
    """
    
    def __init__(self, id: str, max_iterations: int = 10, **config):
        """
        Initialize loop node.
        
        Args:
            id: Node ID
            max_iterations: Maximum number of iterations
            **config: Additional configuration
        """
        config["max_iterations"] = max_iterations
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.LOOP, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="items", type="list", required=True, description="Items to iterate over")
        ]
        self.metadata.outputs = [
            NodeOutput(name="item", type="any", description="Current item"),
            NodeOutput(name="index", type="int", description="Current index"),
            NodeOutput(name="results", type="list", description="All results")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Iterate over items.
        
        Note: This is a simplified implementation. Full loop support would require
        workflow engine modifications to handle iterative execution.
        
        Args:
            inputs: Items to iterate
            context: Execution context
        
        Returns:
            Iteration results
        """
        items = inputs.get("items", [])
        max_iterations = self.metadata.config.get("max_iterations", 10)
        
        if not isinstance(items, list):
            items = [items]
        
        # Limit iterations
        items = items[:max_iterations]
        
        # For now, return all items (full loop execution would need engine support)
        return {
            "results": items,
            "item": items[-1] if items else None,
            "index": len(items) - 1 if items else 0
        }


class MergeNode(BaseNode):
    """
    Merge node to combine multiple inputs.
    """
    
    def __init__(self, id: str, merge_strategy: str = "list", **config):
        """
        Initialize merge node.
        
        Args:
            id: Node ID
            merge_strategy: 'list', 'dict', or 'concat'
            **config: Additional configuration
        """
        config["merge_strategy"] = merge_strategy
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.MERGE, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="input1", type="any", required=False, description="First input"),
            NodeInput(name="input2", type="any", required=False, description="Second input"),
            NodeInput(name="input3", type="any", required=False, description="Third input"),
        ]
        self.metadata.outputs = [
            NodeOutput(name="output", type="any", description="Merged output")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Merge multiple inputs.
        
        Args:
            inputs: Multiple input values
            context: Execution context
        
        Returns:
            Merged output
        """
        strategy = self.metadata.config.get("merge_strategy", "list")
        
        # Collect all non-None inputs
        values = [v for k, v in inputs.items() if v is not None and k.startswith("input")]
        
        if strategy == "list":
            result = values
        elif strategy == "dict":
            result = {f"input{i+1}": v for i, v in enumerate(values)}
        elif strategy == "concat":
            # Concatenate strings or lists
            result = ""
            for v in values:
                if isinstance(v, str):
                    result += v
                elif isinstance(v, list):
                    if isinstance(result, str):
                        result = []
                    result.extend(v)
                else:
                    result += str(v)
        else:
            result = values
        
        return {"output": result}
