"""
Execution Engine for NeuroFlow

Handles topological sorting of the DAG and orchestrates
node execution in the correct order.
"""

from __future__ import annotations

import asyncio
import time
from datetime import datetime
from typing import Any, Callable, Optional
from uuid import uuid4

import networkx as nx

from .schemas import (
    Edge,
    ExecutionProgress,
    ExecutionResult,
    ExecutionStatus,
    Node,
    NodeExecutionResult,
    Workflow,
)
from .storage import StorageManager


class DAGValidationError(Exception):
    """Raised when the workflow DAG is invalid."""
    pass


class NodeExecutionError(Exception):
    """Raised when a node fails to execute."""
    pass


class ExecutionEngine:
    """
    Core execution engine for NeuroFlow workflows.
    
    Responsibilities:
    - Validate workflow DAG structure
    - Perform topological sort to determine execution order
    - Execute nodes in order, passing data between them
    - Track execution progress and results
    """
    
    def __init__(self, storage: StorageManager):
        self.storage = storage
        self._executors: dict[str, Callable] = {}
        self._running_tasks: dict[str, asyncio.Task] = {}
        self._progress_callbacks: dict[str, Callable] = {}
        
    def register_executor(self, node_type: str, executor: Callable) -> None:
        """Register an executor function for a node type."""
        self._executors[node_type] = executor
        
    def get_executor(self, node_type: str) -> Optional[Callable]:
        """Get the executor for a node type."""
        return self._executors.get(node_type)
    
    def build_graph(self, workflow: Workflow) -> nx.DiGraph:
        """
        Build a NetworkX directed graph from the workflow.
        
        Args:
            workflow: The workflow to convert
            
        Returns:
            NetworkX DiGraph representing the workflow
        """
        graph = nx.DiGraph()
        
        # Add all nodes
        for node in workflow.nodes:
            graph.add_node(node.id, node=node)
            
        # Add all edges
        for edge in workflow.edges:
            graph.add_edge(edge.source, edge.target, edge=edge)
            
        return graph
    
    def validate_dag(self, workflow: Workflow) -> tuple[bool, str]:
        """
        Validate that the workflow forms a valid DAG.
        
        Checks:
        - No cycles exist
        - All edge references point to valid nodes
        - Graph is connected (optional warning)
        
        Args:
            workflow: The workflow to validate
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        graph = self.build_graph(workflow)
        
        # Check for cycles
        if not nx.is_directed_acyclic_graph(graph):
            cycles = list(nx.simple_cycles(graph))
            return False, f"Workflow contains cycles: {cycles}"
        
        # Check that all edge endpoints exist
        node_ids = {node.id for node in workflow.nodes}
        for edge in workflow.edges:
            if edge.source not in node_ids:
                return False, f"Edge source '{edge.source}' does not exist"
            if edge.target not in node_ids:
                return False, f"Edge target '{edge.target}' does not exist"
        
        return True, ""
    
    def topological_sort(self, workflow: Workflow) -> list[str]:
        """
        Perform topological sort on the workflow DAG.
        
        Args:
            workflow: The workflow to sort
            
        Returns:
            List of node IDs in execution order
            
        Raises:
            DAGValidationError: If the graph is not a valid DAG
        """
        is_valid, error = self.validate_dag(workflow)
        if not is_valid:
            raise DAGValidationError(error)
            
        graph = self.build_graph(workflow)
        return list(nx.topological_sort(graph))
    
    def get_node_by_id(self, workflow: Workflow, node_id: str) -> Optional[Node]:
        """Get a node from the workflow by its ID."""
        for node in workflow.nodes:
            if node.id == node_id:
                return node
        return None
    
    def get_input_node_ids(self, workflow: Workflow, node_id: str) -> list[str]:
        """Get the IDs of all nodes that feed into the given node."""
        return [edge.source for edge in workflow.edges if edge.target == node_id]
    
    async def execute_node(
        self,
        node: Node,
        input_data: dict[str, Any],
        run_id: str,
    ) -> tuple[Any, NodeExecutionResult]:
        """
        Execute a single node.
        
        Args:
            node: The node to execute
            input_data: Data from upstream nodes
            run_id: Current run ID for storage
            
        Returns:
            Tuple of (output_data, execution_result)
        """
        start_time = time.perf_counter()
        
        # Get the executor for this node type
        executor = self.get_executor(node.type)
        if executor is None:
            raise NodeExecutionError(f"No executor registered for node type: {node.type}")
        
        try:
            # Execute the node
            output_data = await asyncio.to_thread(
                executor,
                node=node,
                input_data=input_data,
                storage=self.storage,
                run_id=run_id,
            )
            
            execution_time = (time.perf_counter() - start_time) * 1000  # ms
            
            # Create output summary (limited info for large DataFrames)
            output_summary = self._create_output_summary(output_data)
            
            result = NodeExecutionResult(
                node_id=node.id,
                status=ExecutionStatus.COMPLETED,
                execution_time_ms=execution_time,
                output_summary=output_summary,
            )
            
            return output_data, result
            
        except Exception as e:
            execution_time = (time.perf_counter() - start_time) * 1000
            
            result = NodeExecutionResult(
                node_id=node.id,
                status=ExecutionStatus.FAILED,
                execution_time_ms=execution_time,
                error_message=str(e),
            )
            
            raise NodeExecutionError(f"Node {node.id} failed: {e}") from e
    
    def _create_output_summary(self, output_data: Any, is_final: bool = False) -> dict[str, Any]:
        """Create a summary of node output for logging."""
        import pandas as pd
        
        summary = {"type": type(output_data).__name__}
        
        if isinstance(output_data, pd.DataFrame):
            summary.update({
                "shape": output_data.shape,
                "columns": list(output_data.columns)[:10],
                "dtypes": {str(k): str(v) for k, v in output_data.dtypes.items()},
            })
        elif isinstance(output_data, dict):
            summary["keys"] = list(output_data.keys())
            # If this is the final output and contains metrics, include them directly
            if is_final:
                if "metrics" in output_data:
                    summary["metrics"] = output_data["metrics"]
                if "cross_validation" in output_data:
                    summary["cross_validation"] = output_data["cross_validation"]
                if "feature_importances" in output_data:
                    summary["feature_importances"] = output_data["feature_importances"]
                if "model_type" in output_data:
                    summary["model_type"] = output_data["model_type"]
                if "is_classifier" in output_data:
                    summary["is_classifier"] = output_data["is_classifier"]
                if "scatter_data" in output_data:
                    summary["scatter_data"] = output_data["scatter_data"]
                if "qq_data" in output_data:
                    summary["qq_data"] = output_data["qq_data"]
                if "residuals" in output_data:
                    summary["residuals"] = output_data["residuals"][:100]  # Limit size
                if "y_test" in output_data:
                    summary["y_test"] = output_data["y_test"][:100]
                if "y_pred" in output_data:
                    summary["y_pred"] = output_data["y_pred"][:100]
                if "confusion_matrix" in output_data:
                    summary["confusion_matrix"] = output_data["confusion_matrix"]
                if "model_path" in output_data:
                    summary["model_path"] = output_data["model_path"]
        elif hasattr(output_data, "__len__"):
            summary["length"] = len(output_data)
            
        return summary
    
    async def execute_workflow(
        self,
        workflow: Workflow,
        run_id: Optional[str] = None,
        progress_callback: Optional[Callable[[ExecutionProgress], None]] = None,
    ) -> ExecutionResult:
        """
        Execute a complete workflow.
        
        Args:
            workflow: The workflow to execute
            run_id: Optional run ID (generated if not provided)
            progress_callback: Optional callback for progress updates
            
        Returns:
            ExecutionResult with all node results
        """
        run_id = run_id or str(uuid4())
        started_at = datetime.utcnow()
        
        # Initialize result
        result = ExecutionResult(
            run_id=run_id,
            workflow_id=workflow.id,
            status=ExecutionStatus.RUNNING,
            started_at=started_at,
            node_results=[],
        )
        
        try:
            # Validate and sort
            execution_order = self.topological_sort(workflow)
            total_nodes = len(execution_order)
            
            # Storage for intermediate data
            node_outputs: dict[str, Any] = {}
            
            # Execute nodes in order
            for index, node_id in enumerate(execution_order):
                node = self.get_node_by_id(workflow, node_id)
                if node is None:
                    raise NodeExecutionError(f"Node not found: {node_id}")
                
                # Send progress update
                if progress_callback:
                    progress = ExecutionProgress(
                        run_id=run_id,
                        workflow_id=workflow.id,
                        current_node_id=node_id,
                        current_node_index=index,
                        total_nodes=total_nodes,
                        status=ExecutionStatus.RUNNING,
                        message=f"Executing node: {node.data.label or node.type}",
                    )
                    await asyncio.to_thread(progress_callback, progress)
                
                # Gather input data from upstream nodes
                input_node_ids = self.get_input_node_ids(workflow, node_id)
                input_data = {
                    input_id: node_outputs.get(input_id)
                    for input_id in input_node_ids
                }
                
                # Execute the node
                try:
                    output_data, node_result = await self.execute_node(
                        node=node,
                        input_data=input_data,
                        run_id=run_id,
                    )
                    node_outputs[node_id] = output_data
                    result.node_results.append(node_result)
                    
                except NodeExecutionError as e:
                    node_result = NodeExecutionResult(
                        node_id=node_id,
                        status=ExecutionStatus.FAILED,
                        execution_time_ms=0,
                        error_message=str(e),
                    )
                    result.node_results.append(node_result)
                    raise
            
            # Success
            result.status = ExecutionStatus.COMPLETED
            result.completed_at = datetime.utcnow()
            result.total_execution_time_ms = sum(
                r.execution_time_ms for r in result.node_results
            )
            
            # Get final output (from last node in execution order)
            if execution_order:
                last_node_id = execution_order[-1]
                last_output = node_outputs.get(last_node_id)
                result.final_output = self._create_output_summary(last_output, is_final=True)
            
            # Save execution result
            await asyncio.to_thread(
                self.storage.save_run_result,
                run_id=run_id,
                result=result,
            )
            
            return result
            
        except Exception as e:
            result.status = ExecutionStatus.FAILED
            result.completed_at = datetime.utcnow()
            result.error_message = str(e)
            result.total_execution_time_ms = sum(
                r.execution_time_ms for r in result.node_results
            )
            
            # Save failed result
            await asyncio.to_thread(
                self.storage.save_run_result,
                run_id=run_id,
                result=result,
            )
            
            return result
    
    def cancel_execution(self, run_id: str) -> bool:
        """Cancel a running workflow execution."""
        task = self._running_tasks.get(run_id)
        if task and not task.done():
            task.cancel()
            return True
        return False
