"""
Core workflow engine for YBAgent.

Handles workflow creation, execution, and state management.
"""

from typing import Any, Dict, List, Optional, Set
from enum import Enum
import json
from datetime import datetime
from ybagent.node import BaseNode
from ybagent.edge import Edge, EdgeValidator


class WorkflowStatus(str, Enum):
    """Workflow execution status."""
    IDLE = "idle"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    PAUSED = "paused"


class WorkflowState:
    """
    Manages the state during workflow execution.
    """
    
    def __init__(self):
        self.variables: Dict[str, Any] = {}
        self.node_outputs: Dict[str, Dict[str, Any]] = {}
        self.execution_history: List[Dict[str, Any]] = []
        self.status: WorkflowStatus = WorkflowStatus.IDLE
        self.error: Optional[str] = None
        self.start_time: Optional[datetime] = None
        self.end_time: Optional[datetime] = None
    
    def set_variable(self, name: str, value: Any):
        """Set a workflow variable."""
        self.variables[name] = value
    
    def get_variable(self, name: str, default: Any = None) -> Any:
        """Get a workflow variable."""
        return self.variables.get(name, default)
    
    def set_node_output(self, node_id: str, output: Dict[str, Any]):
        """Store the output of a node."""
        self.node_outputs[node_id] = output
    
    def get_node_output(self, node_id: str) -> Optional[Dict[str, Any]]:
        """Get the output of a node."""
        return self.node_outputs.get(node_id)
    
    def add_execution_record(self, node_id: str, status: str, duration: float, error: Optional[str] = None):
        """Add an execution record to history."""
        self.execution_history.append({
            "node_id": node_id,
            "status": status,
            "duration": duration,
            "timestamp": datetime.now().isoformat(),
            "error": error
        })
    
    def to_dict(self) -> Dict[str, Any]:
        """Serialize state to dictionary."""
        return {
            "variables": self.variables,
            "node_outputs": self.node_outputs,
            "execution_history": self.execution_history,
            "status": self.status.value,
            "error": self.error,
            "start_time": self.start_time.isoformat() if self.start_time else None,
            "end_time": self.end_time.isoformat() if self.end_time else None,
        }


class Workflow:
    """
    Main workflow container that holds nodes and edges.
    """
    
    def __init__(self, name: str = "Untitled Workflow", description: str = ""):
        self.name = name
        self.description = description
        self.nodes: Dict[str, BaseNode] = {}
        self.edges: List[Edge] = []
        self.metadata: Dict[str, Any] = {
            "created_at": datetime.now().isoformat(),
            "version": "1.0"
        }
    
    def add_node(self, node: BaseNode) -> "Workflow":
        """
        Add a node to the workflow.
        
        Args:
            node: Node to add
        
        Returns:
            Self for method chaining
        """
        if node.metadata.id in self.nodes:
            raise ValueError(f"Node with id '{node.metadata.id}' already exists")
        self.nodes[node.metadata.id] = node
        return self
    
    def remove_node(self, node_id: str) -> "Workflow":
        """
        Remove a node from the workflow.
        
        Args:
            node_id: ID of node to remove
        
        Returns:
            Self for method chaining
        """
        if node_id not in self.nodes:
            raise ValueError(f"Node '{node_id}' not found")
        
        # Remove associated edges
        self.edges = [
            e for e in self.edges
            if e.source_node_id != node_id and e.target_node_id != node_id
        ]
        
        del self.nodes[node_id]
        return self
    
    def connect(
        self,
        source_id: str,
        target_id: str,
        source_output: str = "output",
        target_input: str = "input",
        condition: Optional[str] = None
    ) -> "Workflow":
        """
        Connect two nodes with an edge.
        
        Args:
            source_id: Source node ID
            target_id: Target node ID
            source_output: Source output port name
            target_input: Target input port name
            condition: Optional condition for the edge
        
        Returns:
            Self for method chaining
        """
        edge = Edge(
            id=f"{source_id}_{source_output}_to_{target_id}_{target_input}",
            source_node_id=source_id,
            source_output=source_output,
            target_node_id=target_id,
            target_input=target_input,
            condition=condition
        )
        
        # Validate edge
        EdgeValidator.validate(edge, self.nodes)
        
        self.edges.append(edge)
        return self
    
    def disconnect(self, source_id: str, target_id: str) -> "Workflow":
        """
        Remove edge(s) between two nodes.
        
        Args:
            source_id: Source node ID
            target_id: Target node ID
        
        Returns:
            Self for method chaining
        """
        self.edges = [
            e for e in self.edges
            if not (e.source_node_id == source_id and e.target_node_id == target_id)
        ]
        return self
    
    def validate(self) -> bool:
        """
        Validate the workflow structure.
        
        Returns:
            True if workflow is valid
        
        Raises:
            ValueError: If workflow is invalid
        """
        # Check for cycles
        if EdgeValidator.detect_cycles(self.edges, self.nodes):
            raise ValueError("Workflow contains cycles")
        
        # Validate all edges
        for edge in self.edges:
            EdgeValidator.validate(edge, self.nodes)
        
        return True
    
    def execute(self, inputs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Execute the workflow.
        
        Args:
            inputs: Initial input values
        
        Returns:
            Workflow execution results
        """
        executor = WorkflowExecutor(self)
        return executor.execute(inputs or {})
    
    def to_json(self, filepath: Optional[str] = None) -> str:
        """
        Serialize workflow to JSON.
        
        Args:
            filepath: Optional file path to save JSON
        
        Returns:
            JSON string representation
        """
        data = {
            "name": self.name,
            "description": self.description,
            "metadata": self.metadata,
            "nodes": [node.to_dict() for node in self.nodes.values()],
            "edges": [edge.model_dump() for edge in self.edges]
        }
        
        json_str = json.dumps(data, indent=2)
        
        if filepath:
            with open(filepath, 'w') as f:
                f.write(json_str)
        
        return json_str
    
    @classmethod
    def from_json(cls, json_str: Optional[str] = None, filepath: Optional[str] = None) -> "Workflow":
        """
        Deserialize workflow from JSON.
        
        Args:
            json_str: JSON string
            filepath: Optional file path to load JSON from
        
        Returns:
            Workflow instance
        """
        if filepath:
            with open(filepath, 'r') as f:
                json_str = f.read()
        
        if not json_str:
            raise ValueError("Either json_str or filepath must be provided")
        
        data = json.loads(json_str)
        
        workflow = cls(name=data["name"], description=data.get("description", ""))
        workflow.metadata = data.get("metadata", {})
        
        # Import nodes (this requires node registry for custom nodes)
        from ybagent.plugins import NodeRegistry
        for node_data in data["nodes"]:
            node_class = NodeRegistry.get(node_data["type"])
            if node_class:
                node = node_class.from_dict(node_data)
                workflow.add_node(node)
        
        # Import edges
        for edge_data in data["edges"]:
            workflow.edges.append(Edge(**edge_data))
        
        return workflow


class WorkflowExecutor:
    """
    Executes workflows by traversing the node graph.
    """
    
    def __init__(self, workflow: Workflow):
        self.workflow = workflow
        self.state = WorkflowState()
    
    def execute(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute the workflow.
        
        Args:
            inputs: Initial input values
        
        Returns:
            Final workflow outputs
        """
        self.state.status = WorkflowStatus.RUNNING
        self.state.start_time = datetime.now()
        
        try:
            # Validate workflow
            self.workflow.validate()
            
            # Find execution order (topological sort)
            execution_order = self._topological_sort()
            
            # Execute nodes in order
            for node_id in execution_order:
                node = self.workflow.nodes[node_id]
                
                # Gather inputs for this node
                node_inputs = self._gather_node_inputs(node_id, inputs)
                
                # Execute node
                start = datetime.now()
                
                # Retry logic
                max_retries = node.metadata.retry_config.get("max_retries", 0)
                retry_delay = node.metadata.retry_config.get("retry_delay", 1.0)
                retry_on_error = node.metadata.retry_config.get("retry_on_error", True)
                
                retries = 0
                while True:
                    try:
                        node.validate_inputs(node_inputs)
                        output = node.execute(node_inputs, {"state": self.state})
                        self.state.set_node_output(node_id, output)
                        
                        duration = (datetime.now() - start).total_seconds()
                        self.state.add_execution_record(node_id, "success", duration)
                        break # Success, exit retry loop
                        
                    except Exception as e:
                        if retry_on_error and retries < max_retries:
                            retries += 1
                            import time
                            time.sleep(retry_delay)
                            continue
                        
                        # Failed after retries
                        duration = (datetime.now() - start).total_seconds()
                        error_msg = f"Error in node '{node_id}': {str(e)}"
                        self.state.add_execution_record(node_id, "failed", duration, error_msg)
                        raise RuntimeError(error_msg) from e
            
            self.state.status = WorkflowStatus.COMPLETED
            
        except Exception as e:
            self.state.status = WorkflowStatus.FAILED
            self.state.error = str(e)
            raise
        
        finally:
            self.state.end_time = datetime.now()
        
        # Collect outputs from output nodes
        return self._collect_outputs()
    
    def _topological_sort(self) -> List[str]:
        """
        Perform topological sort to determine execution order.
        
        Returns:
            List of node IDs in execution order
        """
        # Build adjacency list and in-degree count
        graph = {node_id: [] for node_id in self.workflow.nodes}
        in_degree = {node_id: 0 for node_id in self.workflow.nodes}
        
        for edge in self.workflow.edges:
            graph[edge.source_node_id].append(edge.target_node_id)
            in_degree[edge.target_node_id] += 1
        
        # Find nodes with no incoming edges (start nodes)
        queue = [node_id for node_id, degree in in_degree.items() if degree == 0]
        result = []
        
        while queue:
            node_id = queue.pop(0)
            result.append(node_id)
            
            for neighbor in graph[node_id]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)
        
        if len(result) != len(self.workflow.nodes):
            raise ValueError("Workflow contains cycles or disconnected nodes")
        
        return result
    
    def _gather_node_inputs(self, node_id: str, initial_inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Gather inputs for a node from connected edges.
        
        Args:
            node_id: Node ID
            initial_inputs: Initial workflow inputs
        
        Returns:
            Dictionary of input values for the node
        """
        node_inputs = {}
        
        # Find incoming edges
        incoming_edges = [e for e in self.workflow.edges if e.target_node_id == node_id]
        
        if not incoming_edges:
            # This is a start node, use initial inputs
            node_inputs = initial_inputs.copy()
        else:
            # Gather outputs from source nodes
            for edge in incoming_edges:
                source_output = self.state.get_node_output(edge.source_node_id)
                if source_output:
                    # Map source output to target input
                    if edge.source_output in source_output:
                        node_inputs[edge.target_input] = source_output[edge.source_output]
                    else:
                        # If specific output not found, pass entire output
                        node_inputs[edge.target_input] = source_output
        
        return node_inputs
    
    def _collect_outputs(self) -> Dict[str, Any]:
        """
        Collect outputs from output nodes.
        
        Returns:
            Dictionary of workflow outputs
        """
        from ybagent.node import NodeType
        
        outputs = {}
        
        # Find output nodes
        for node_id, node in self.workflow.nodes.items():
            if node.metadata.node_type == NodeType.OUTPUT:
                node_output = self.state.get_node_output(node_id)
                if node_output:
                    outputs[node_id] = node_output
        
        # If no output nodes, return all node outputs
        if not outputs:
            outputs = self.state.node_outputs
        
        return outputs
