"""Pipeline run summary and tracking.

This module provides functionality for tracking pipeline execution and generating
JSON summaries of runs.
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from sleap_roots_analyze.data_utils import convert_to_json_serializable


@dataclass
class StepSummary:
    """Summary of a single pipeline step (task) execution.

    Attributes:
        name: Name of the task.
        description: Description of what the task does.
        status: Execution status (success, failed, skipped).
        elapsed_time: Time taken to execute in seconds.
        files_generated: List of file paths generated by this step.
        metadata: Additional metadata about the step execution.
        error: Error message if the step failed.
    """

    name: str
    description: str = ""
    status: str = "pending"
    elapsed_time: float = 0.0
    files_generated: List[str | Path] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None


@dataclass
class PipelineSummary:
    """Summary of an entire pipeline run.

    Attributes:
        pipeline_name: Name of the pipeline.
        version: Pipeline version.
        start_time: ISO format timestamp of when the pipeline started.
        end_time: ISO format timestamp of when the pipeline ended.
        total_elapsed_time: Total time taken for the pipeline in seconds.
        status: Overall pipeline status (success, failed, partial).
        steps: List of StepSummary objects for each task.
        config: Configuration used for this run.
        environment: Environment information (git commit, package versions, etc.).
        output_directory: Directory where outputs were saved.
        data_source: Input data path(s) used for this run.
    """

    pipeline_name: str
    version: str = "1.0"
    start_time: str = ""
    end_time: str = ""
    total_elapsed_time: float = 0.0
    status: str = "running"
    steps: List[StepSummary] = field(default_factory=list)
    config: Dict[str, Any] = field(default_factory=dict)
    environment: Dict[str, Any] = field(default_factory=dict)
    output_directory: str = ""
    data_source: str = ""

    def to_dict(self) -> Dict[str, Any]:
        """Convert summary to dictionary.

        Returns:
            Dictionary representation of the summary.
        """
        return asdict(self)

    def to_json(self, indent: int = 2) -> str:
        """Convert summary to JSON string.

        Args:
            indent: Number of spaces for JSON indentation.

        Returns:
            JSON string representation of the summary.
        """
        # Convert Path objects to strings and numpy types to native Python
        data = self.to_dict()
        data["steps"] = [
            {
                **step,
                "files_generated": [str(f) for f in step["files_generated"]],
            }
            for step in data["steps"]
        ]
        # Use convert_to_json_serializable to handle numpy types in metadata
        data = convert_to_json_serializable(data)
        return json.dumps(data, indent=indent)

    def save(self, path: Path | str) -> None:
        """Save summary to a JSON file.

        Args:
            path: Path to save the JSON file.
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(self.to_json())

    @classmethod
    def load(cls, path: Path | str) -> PipelineSummary:
        """Load summary from a JSON file.

        Args:
            path: Path to the JSON file.

        Returns:
            PipelineSummary object loaded from the file.
        """
        path = Path(path)
        data = json.loads(path.read_text())

        # Convert step dictionaries back to StepSummary objects
        steps = [StepSummary(**step) for step in data.pop("steps", [])]

        return cls(steps=steps, **data)

    def mark_step_success(
        self,
        step_name: str,
        elapsed_time: float,
        files_generated: List[str | Path],
        metadata: Dict[str, Any],
    ) -> None:
        """Mark a step as successfully completed.

        Args:
            step_name: Name of the step.
            elapsed_time: Time taken in seconds.
            files_generated: List of files generated.
            metadata: Additional metadata.
        """
        for step in self.steps:
            if step.name == step_name:
                step.status = "success"
                step.elapsed_time = elapsed_time
                step.files_generated = files_generated
                step.metadata = metadata
                return

    def mark_step_failed(self, step_name: str, error: str) -> None:
        """Mark a step as failed.

        Args:
            step_name: Name of the step.
            error: Error message.
        """
        for step in self.steps:
            if step.name == step_name:
                step.status = "failed"
                step.error = error
                return

    def finalize(self, status: str = "success") -> None:
        """Finalize the pipeline summary.

        Args:
            status: Overall pipeline status (success, failed, partial).
        """
        self.end_time = datetime.now().isoformat()
        self.status = status

        # Calculate total elapsed time if start_time is set
        if self.start_time:
            start = datetime.fromisoformat(self.start_time)
            end = datetime.fromisoformat(self.end_time)
            self.total_elapsed_time = (end - start).total_seconds()
