"""Workflow visualization utilities."""

import matplotlib

matplotlib.use("Agg")  # Use non-interactive backend
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import networkx as nx

from kailash.tracking.manager import TaskManager
from kailash.workflow.graph import Workflow


class WorkflowVisualizer:
    """Visualize workflows as graphs."""

    def __init__(
        self,
        workflow: Workflow,
        node_colors: dict[str, str] | None = None,
        edge_colors: dict[str, str] | None = None,
        layout: str = "hierarchical",
    ):
        """Initialize visualizer.

        Args:
            workflow: Workflow to visualize
            node_colors: Custom node color map
            edge_colors: Custom edge color map
            layout: Layout algorithm to use
        """
        self.workflow = workflow
        self.node_colors = node_colors or self._default_node_colors()
        self.edge_colors = edge_colors or self._default_edge_colors()
        self.layout = layout

    def _default_node_colors(self) -> dict[str, str]:
        """Get default node color map."""
        return {
            "data": "lightblue",
            "transform": "lightyellow",
            "logic": "lightcoral",
            "ai": "lightpink",
            "default": "lightgray",
        }

    def _default_edge_colors(self) -> dict[str, str]:
        """Get default edge color map."""
        return {"default": "gray", "error": "red", "conditional": "orange"}

    def _get_node_color(self, node_type: str) -> str:
        """Get color for node based on type."""
        if "Reader" in node_type or "Writer" in node_type:
            return self.node_colors["data"]
        elif (
            "Transform" in node_type
            or "Filter" in node_type
            or "Processor" in node_type
        ):
            return self.node_colors["transform"]
        elif "Logic" in node_type or "Merge" in node_type or "Conditional" in node_type:
            return self.node_colors["logic"]
        elif "AI" in node_type or "Model" in node_type:
            return self.node_colors["ai"]
        return self.node_colors["default"]

    def _get_node_colors(self) -> list[str]:
        """Get colors for all nodes in workflow."""
        colors = []
        for node_id in self.workflow.graph.nodes():
            node_instance = self.workflow.nodes[node_id]
            node_type = node_instance.node_type
            colors.append(self._get_node_color(node_type))
        return colors

    def _get_node_labels(self) -> dict[str, str]:
        """Get labels for nodes in workflow."""
        labels = {}
        for node_id in self.workflow.graph.nodes():
            # Try to get name from node instance
            node = self.workflow.get_node(node_id)
            if node and hasattr(node, "name") and node.name:
                # For test compatibility - test expects just the name, not node_id
                labels[node_id] = node.name
            else:
                # Fallback to node type from metadata
                node_instance = self.workflow.nodes.get(node_id)
                if node_instance:
                    labels[node_id] = f"{node_id} ({node_instance.node_type})"
                else:
                    labels[node_id] = node_id
        return labels

    def _get_edge_labels(self) -> dict[tuple[str, str], str]:
        """Get labels for edges in workflow."""
        edge_labels = {}

        for edge in self.workflow.graph.edges(data=True):
            source, target, data = edge

            # Try both edge data formats for compatibility
            from_output = data.get("from_output")
            to_input = data.get("to_input")

            if from_output and to_input:
                edge_labels[(source, target)] = f"{from_output}→{to_input}"
            else:
                # Fallback to mapping format
                mapping = data.get("mapping", {})
                if mapping:
                    # Create label from mapping
                    label_parts = []
                    for src, dst in mapping.items():
                        label_parts.append(f"{src}→{dst}")
                    label = "\n".join(label_parts)
                    edge_labels[(source, target)] = label

        return edge_labels

    def _calculate_layout(self) -> dict[str, tuple[float, float]]:
        """Calculate node positions for visualization."""
        # Try to use stored positions first
        pos = {}
        for node_id, node_instance in self.workflow.nodes.items():
            if node_instance.position != (0, 0):
                pos[node_id] = node_instance.position

        # If no positions stored, calculate them
        if not pos:
            if self.layout == "hierarchical":
                # Use hierarchical layout for DAGs
                try:
                    # Create layers based on topological order
                    layers = self._create_layers()
                    pos = self._hierarchical_layout(layers)
                except Exception:
                    # Fallback to spring layout
                    pos = nx.spring_layout(self.workflow.graph, k=3, iterations=50)
            elif self.layout == "circular":
                pos = nx.circular_layout(self.workflow.graph)
            elif self.layout == "spring":
                pos = nx.spring_layout(self.workflow.graph, k=2, iterations=100)
            else:
                # Default to spring layout
                pos = nx.spring_layout(self.workflow.graph)

        return pos

    def _create_layers(self) -> dict[int, list]:
        """Create layers of nodes for hierarchical layout."""
        layers = {}
        remaining = set(self.workflow.graph.nodes())
        layer = 0

        while remaining:
            # Find nodes with no dependencies in remaining set
            current_layer = []
            for node in remaining:
                predecessors = set(self.workflow.graph.predecessors(node))
                if not predecessors.intersection(remaining):
                    current_layer.append(node)

            if not current_layer:
                # Circular dependency, break and use all remaining
                current_layer = list(remaining)

            layers[layer] = current_layer
            remaining -= set(current_layer)
            layer += 1

        return layers

    def _hierarchical_layout(
        self, layers: dict[int, list]
    ) -> dict[str, tuple[float, float]]:
        """Create hierarchical layout from layers."""
        pos = {}
        layer_height = 2.0

        for layer, nodes in layers.items():
            y = layer * layer_height
            if len(nodes) == 1:
                x_positions = [0]
            else:
                width = max(2.0, len(nodes) - 1)
                x_positions = [
                    -width / 2 + i * width / (len(nodes) - 1) for i in range(len(nodes))
                ]

            for i, node in enumerate(nodes):
                pos[node] = (x_positions[i], -y)  # Negative y to flow top to bottom

        return pos

    def _draw_graph(
        self,
        pos: dict[str, tuple[float, float]],
        node_colors: list[str],
        show_labels: bool,
        show_connections: bool,
    ) -> None:
        """Draw the graph with given positions and options."""
        # Draw nodes
        nx.draw_networkx_nodes(
            self.workflow.graph, pos, node_color=node_colors, node_size=3000, alpha=0.9
        )

        # Draw edges
        nx.draw_networkx_edges(
            self.workflow.graph,
            pos,
            edge_color=self.edge_colors["default"],
            width=2,
            alpha=0.6,
            arrows=True,
            arrowsize=20,
            arrowstyle="->",
        )

        # Draw labels
        if show_labels:
            labels = self._get_node_labels()
            nx.draw_networkx_labels(
                self.workflow.graph, pos, labels, font_size=10, font_weight="bold"
            )

        # Draw connection labels
        if show_connections:
            edge_labels = self._get_edge_labels()
            nx.draw_networkx_edge_labels(
                self.workflow.graph, pos, edge_labels, font_size=8
            )

    def visualize(
        self,
        output_path: str | None = None,
        figsize: tuple[int, int] = (12, 8),
        title: str | None = None,
        show_labels: bool = True,
        show_connections: bool = True,
        dpi: int = 300,
        **kwargs,
    ) -> None:
        """Create a visualization of the workflow.

        Args:
            output_path: Path to save the image (if None, shows interactive plot)
            figsize: Figure size (width, height)
            title: Optional title for the graph
            show_labels: Whether to show node labels
            show_connections: Whether to show connection labels
            dpi: Resolution (dots per inch) for saved images
            **kwargs: Additional options passed to plt.savefig
        """
        try:
            plt.figure(figsize=figsize)

            # Calculate node positions
            pos = self._calculate_layout()

            # Handle empty workflow case
            if not self.workflow.graph.nodes():
                pos = {}
                node_colors = []
            else:
                # Draw the graph with colors
                node_colors = self._get_node_colors()

            # Draw the graph components
            if pos and node_colors:
                self._draw_graph(pos, node_colors, show_labels, show_connections)

            # Set title
            if title is None:
                title = f"Workflow: {self.workflow.name}"
            plt.title(title, fontsize=16, fontweight="bold")

            # Remove axes
            plt.axis("off")
            plt.tight_layout()

            # Show or save
            if output_path:
                plt.savefig(output_path, dpi=dpi, bbox_inches="tight", **kwargs)
                plt.close()
            else:
                plt.show()
        except Exception as e:
            plt.close()
            raise e

    def save(self, output_path: str, dpi: int = 300, **kwargs) -> None:
        """Save visualization to file.

        Args:
            output_path: Path to save the image
            dpi: Resolution (dots per inch)
            **kwargs: Additional options for plt.savefig
        """
        kwargs["dpi"] = dpi
        self.visualize(output_path=output_path, **kwargs)

    def create_execution_graph(
        self, run_id: str, task_manager: Any, output_path: str | None = None
    ) -> str:
        """Create a Mermaid visualization showing execution status.

        Args:
            run_id: Run ID to visualize
            task_manager: Task manager instance
            output_path: Optional path to save the markdown file. If not provided,
                        saves to 'workflow_executions/execution_{run_id}.md'

        Returns:
            Path to the created markdown file
        """
        # Import here to avoid circular dependency
        from kailash.tracking import TaskStatus
        from kailash.workflow.mermaid_visualizer import MermaidVisualizer

        # Get tasks for this run
        tasks = task_manager.list_tasks(run_id)

        # Create status emoji mapping
        status_emojis = {
            TaskStatus.PENDING: "⏳",
            TaskStatus.RUNNING: "🔄",
            TaskStatus.COMPLETED: "✅",
            TaskStatus.FAILED: "❌",
            TaskStatus.SKIPPED: "⏭️",
        }

        # Map node IDs to statuses
        node_status = {}
        for task in tasks:
            node_status[task.node_id] = task.status

        # Create Mermaid visualizer
        visualizer = MermaidVisualizer(self.workflow)

        # Generate base Mermaid diagram
        mermaid_code = visualizer.generate()

        # Add status information to the diagram
        lines = mermaid_code.split("\n")
        new_lines = []

        for line in lines:
            new_lines.append(line)
            # Add status emoji to node labels
            for node_id, status in node_status.items():
                sanitized_id = visualizer._sanitize_node_id(node_id)
                if sanitized_id in line and '["' in line:
                    # Add emoji to the end of the label
                    emoji = status_emojis.get(status, "")
                    if emoji:
                        new_lines[-1] = line.replace('"]', f' {emoji}"]')

        # Create markdown content
        newline_joined = "\n".join(new_lines)
        markdown_content = f"""# Workflow Execution Status

**Run ID**: `{run_id}`
**Workflow**: {self.workflow.name}
**Timestamp**: {task_manager.get_run(run_id).started_at if hasattr(task_manager, 'get_run') else 'N/A'}

## Execution Diagram

```mermaid
{newline_joined}
```

## Status Legend

| Status | Symbol | Description |
|--------|--------|-------------|
| Pending | ⏳ | Task is waiting to be executed |
| Running | 🔄 | Task is currently executing |
| Completed | ✅ | Task completed successfully |
| Failed | ❌ | Task failed during execution |
| Skipped | ⏭️ | Task was skipped |

## Task Details

| Node ID | Status | Start Time | End Time | Duration |
|---------|--------|------------|----------|----------|
"""

        # Add task details
        for task in tasks:
            status_emoji = status_emojis.get(task.status, "")
            start_time = task.started_at or "N/A"
            end_time = task.ended_at or "N/A"
            duration = f"{task.duration:.2f}s" if task.duration else "N/A"

            markdown_content += f"| {task.node_id} | {task.status.value} {status_emoji} | {start_time} | {end_time} | {duration} |\n"

        # Determine output path
        if output_path is None:
            # Use centralized output directory
            # Get project root and use data/outputs/visualizations
            project_root = Path(__file__).parent.parent.parent.parent
            output_dir = (
                project_root
                / "data"
                / "outputs"
                / "visualizations"
                / "workflow_executions"
            )
            output_dir.mkdir(parents=True, exist_ok=True)
            output_path = output_dir / f"execution_{run_id}.md"
        else:
            output_path = Path(output_path)
            # Create parent directory if needed
            output_path.parent.mkdir(parents=True, exist_ok=True)

        # Write markdown file
        with open(output_path, "w") as f:
            f.write(markdown_content)

        return str(output_path)

    def create_performance_dashboard(
        self, run_id: str, task_manager: TaskManager, output_dir: Path | None = None
    ) -> dict[str, Path]:
        """Create integrated performance dashboard with workflow visualization.

        Args:
            run_id: Run ID to visualize
            task_manager: Task manager with execution data
            output_dir: Directory for output files

        Returns:
            Dictionary of output file paths
        """
        from kailash.visualization.performance import PerformanceVisualizer

        if output_dir is None:
            output_dir = Path.cwd() / "outputs" / "performance" / run_id
        output_dir.mkdir(parents=True, exist_ok=True)

        outputs = {}

        # Create workflow execution graph
        outputs["workflow_graph"] = self.create_execution_graph(
            run_id, task_manager, str(output_dir / "workflow_execution.md")
        )

        # Create performance visualizations
        perf_viz = PerformanceVisualizer(task_manager)
        perf_outputs = perf_viz.create_run_performance_summary(run_id, output_dir)
        outputs.update(perf_outputs)

        # Create integrated dashboard HTML
        dashboard_path = output_dir / "dashboard.html"
        self._create_dashboard_html(run_id, outputs, dashboard_path)
        outputs["dashboard"] = dashboard_path

        return outputs

    def _create_dashboard_html(
        self, run_id: str, outputs: dict[str, Path], dashboard_path: Path
    ) -> None:
        """Create HTML dashboard integrating all visualizations."""
        html_content = f"""
<!DOCTYPE html>
<html>
<head>
    <title>Performance Dashboard - Run {run_id}</title>
    <style>
        body {{
            font-family: Arial, sans-serif;
            margin: 20px;
            background-color: #f5f5f5;
        }}
        .container {{
            max-width: 1200px;
            margin: 0 auto;
            background-color: white;
            padding: 20px;
            border-radius: 8px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }}
        h1, h2 {{
            color: #333;
        }}
        .section {{
            margin: 20px 0;
            padding: 20px;
            border: 1px solid #ddd;
            border-radius: 4px;
        }}
        .image-grid {{
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(500px, 1fr));
            gap: 20px;
            margin: 20px 0;
        }}
        img {{
            max-width: 100%;
            height: auto;
            border: 1px solid #ddd;
            border-radius: 4px;
        }}
        .report-link {{
            display: inline-block;
            margin: 10px 0;
            padding: 10px 20px;
            background-color: #007bff;
            color: white;
            text-decoration: none;
            border-radius: 4px;
        }}
        .report-link:hover {{
            background-color: #0056b3;
        }}
    </style>
</head>
<body>
    <div class="container">
        <h1>Performance Dashboard - Run {run_id}</h1>

        <div class="section">
            <h2>Workflow Execution</h2>
            <p>View the workflow execution graph and task status:</p>
            <a href="{outputs.get('workflow_graph', '#')}" class="report-link">
                View Workflow Graph
            </a>
        </div>

        <div class="section">
            <h2>Performance Metrics</h2>
            <div class="image-grid">
"""

        # Add performance images
        image_keys = [
            "execution_timeline",
            "resource_usage",
            "performance_comparison",
            "io_analysis",
            "performance_heatmap",
        ]

        for key in image_keys:
            if key in outputs:
                title = key.replace("_", " ").title()
                html_content += f"""
                <div>
                    <h3>{title}</h3>
                    <img src="{outputs[key].name}" alt="{title}">
                </div>
"""

        html_content += """
            </div>
        </div>

        <div class="section">
            <h2>Detailed Report</h2>
            <p>View the comprehensive performance analysis report:</p>
            <a href="{outputs.get('report', '#')}" class="report-link">
                View Performance Report
            </a>
        </div>
    </div>
</body>
</html>
"""

        with open(dashboard_path, "w") as f:
            f.write(html_content)


# Add visualization method to Workflow class
def add_visualization_to_workflow():
    """Add visualization method to Workflow class."""

    def visualize(self, output_path: str | None = None, **kwargs) -> None:
        """Visualize the workflow.

        Args:
            output_path: Path to save the visualization
            **kwargs: Additional arguments for the visualizer
        """
        visualizer = WorkflowVisualizer(self)
        visualizer.visualize(output_path, **kwargs)

    # Add method to Workflow class
    Workflow.visualize = visualize


# Call this when module is imported
add_visualization_to_workflow()
