"""
Input/Output nodes for workflow data flow.
"""

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


class InputNode(BaseNode):
    """
    Input node for receiving workflow parameters.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.INPUT, **config)

    def _define_inputs_outputs(self):
        # Input nodes don't have inputs, they provide outputs
        self.metadata.outputs = [
            NodeOutput(name="output", type="any", description="Input data")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Pass through input data.
        
        Args:
            inputs: Input data from workflow execution
            context: Execution context
        
        Returns:
            Input data as output
        """
        # If inputs contains a key matching our ID, return that value
        if self.metadata.id in inputs:
            return {"output": inputs[self.metadata.id]}
        
        # Fallback: return the whole inputs dict (legacy behavior)
        return {"output": inputs}


class OutputNode(BaseNode):
    """
    Output node for returning workflow results.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.OUTPUT, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="input", type="any", required=True, description="Data to output")
        ]
        self.metadata.outputs = [
            NodeOutput(name="output", type="any", description="Output data")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Return input data as workflow output.
        
        Args:
            inputs: Input data
            context: Execution context
        
        Returns:
            Input data as output
        """
        return {"output": inputs.get("input")}


class VariableNode(BaseNode):
    """
    Variable node for storing and retrieving workflow variables.
    """
    
    def __init__(self, id: str, variable_name: str = "", operation: str = "get", **config):
        """
        Initialize variable node.
        
        Args:
            id: Node ID
            variable_name: Name of the variable
            operation: 'get' or 'set'
            **config: Additional configuration
        """
        config.pop('node_type', None)
        config["variable_name"] = variable_name
        config["operation"] = operation
        super().__init__(id=id, node_type=NodeType.VARIABLE, **config)
    
    def _define_inputs_outputs(self):
        operation = self.metadata.config.get("operation", "get")
        
        if operation == "set":
            self.metadata.inputs = [
                NodeInput(name="value", type="any", required=True, description="Value to store")
            ]
        
        self.metadata.outputs = [
            NodeOutput(name="output", type="any", description="Variable value")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Get or set a workflow variable.
        
        Args:
            inputs: Input data
            context: Execution context with state
        
        Returns:
            Variable value
        """
        state = context.get("state")
        variable_name = self.metadata.config.get("variable_name", "")
        operation = self.metadata.config.get("operation", "get")
        
        if not variable_name:
            raise ValueError("Variable name not specified")
        
        if operation == "set":
            value = inputs.get("value")
            if state:
                state.set_variable(variable_name, value)
            return {"output": value}
        else:  # get
            value = state.get_variable(variable_name) if state else None
            return {"output": value}
