"""Execution tracking and management for workflow runs."""

import uuid
from datetime import datetime
from typing import Any, Dict, List
from enum import Enum
from .persistence import DatabaseManager, _make_serializable


class ExecutionStatus(str, Enum):
    """Status of a workflow execution."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class Execution:
    """Represents a single workflow execution."""

    def __init__(
        self,
        workflow_name: str,
        inputs: Dict[str, Any],
        execution_id: str | None = None
    ):
        self.execution_id = execution_id or str(uuid.uuid4())
        self.workflow_name = workflow_name
        self.inputs = inputs
        self.status = ExecutionStatus.PENDING
        self.started_at: datetime | None = None
        self.completed_at: datetime | None = None
        self.error: str | None = None
        self.outputs: Any = None
        self.checkpoints: List[str] = []  # List of checkpoint IDs
        self.step_executions: List[Dict[str, Any]] = []  # Track step execution order

        # Replay mode for time-travel: skip cached steps until target
        self.replay_mode: bool = False
        self.replay_from_step: str | None = None
        self.replay_checkpoint_cache: Dict[str, Any] = {}

    def start(self, db: DatabaseManager | None = None):
        """Mark execution as started."""
        self.status = ExecutionStatus.RUNNING
        self.started_at = datetime.now()
        if db:
            self._save_to_db(db)

    def complete(self, outputs: Any, db: DatabaseManager | None = None):
        """Mark execution as completed."""
        self.status = ExecutionStatus.COMPLETED
        self.completed_at = datetime.now()
        self.outputs = outputs
        if db:
            self._save_to_db(db)

    def fail(self, error: str, db: DatabaseManager | None = None):
        """Mark execution as failed."""
        self.status = ExecutionStatus.FAILED
        self.completed_at = datetime.now()
        self.error = error
        if db:
            self._save_to_db(db)

    def _save_to_db(self, db: DatabaseManager):
        """Save execution to database."""
        db.save_execution({
            'execution_id': self.execution_id,
            'workflow_name': self.workflow_name,
            'status': self.status.value,
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None,
            'error': self.error,
            'inputs': self.inputs,
            'outputs': self.outputs
        })

    def add_checkpoint(self, checkpoint_id: str):
        """Add a checkpoint ID to this execution."""
        self.checkpoints.append(checkpoint_id)

    def add_step_execution(self, step_name: str, inputs: Any, outputs: Any, error: str | None = None):
        """Track a step execution.

        Args:
            step_name: Name of the step that was executed
            inputs: Inputs to the step
            outputs: Outputs from the step
            error: Error message if step failed (optional)
        """
        self.step_executions.append({
            "step_name": step_name,
            "inputs": _make_serializable(inputs),
            "outputs": _make_serializable(outputs),
            "error": error,
            "timestamp": datetime.now().isoformat()
        })

    def to_dict(self) -> Dict[str, Any]:
        """Convert execution to dictionary."""
        return {
            "execution_id": self.execution_id,
            "workflow_name": self.workflow_name,
            "status": self.status.value,
            "inputs": _make_serializable(self.inputs),
            "outputs": _make_serializable(self.outputs),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "error": self.error,
            "checkpoints": self.checkpoints,
            "step_executions": self.step_executions,
        }


class ExecutionManager:
    """Manages workflow executions."""

    _instance_counter = 0

    def __init__(self, db_path: str = ".ombra_executions.db"):
        ExecutionManager._instance_counter += 1
        self._instance_id = ExecutionManager._instance_counter
        self._executions: Dict[str, Execution] = {}
        self._current_execution_id: str | None = None
        self.db = DatabaseManager(db_path)

        # Load recent executions from database
        self._load_recent_executions()

    def _load_recent_executions(self, limit: int = 100):
        """Load recent executions from database into memory."""
        executions_data = self.db.load_all_executions(limit=limit)

        for exec_data in executions_data:
            execution = Execution(
                workflow_name=exec_data['workflow_name'],
                inputs=exec_data['inputs'],
                execution_id=exec_data['execution_id']
            )
            execution.status = ExecutionStatus(exec_data['status'])

            if exec_data['started_at']:
                execution.started_at = datetime.fromisoformat(exec_data['started_at'])
            if exec_data['completed_at']:
                execution.completed_at = datetime.fromisoformat(exec_data['completed_at'])

            execution.error = exec_data['error']
            execution.outputs = exec_data['outputs']

            # Load step executions from database
            step_execs = self.db.load_step_executions(execution.execution_id)
            for step_exec in step_execs:
                # Load checkpoint to get inputs/outputs
                if step_exec['checkpoint_id']:
                    checkpoint = self.db.load_checkpoint(step_exec['checkpoint_id'])
                    if checkpoint:
                        execution.step_executions.append({
                            'step_name': step_exec['step_name'],
                            'inputs': checkpoint['inputs'],
                            'outputs': checkpoint['outputs'],
                            'timestamp': step_exec['timestamp']
                        })
                        execution.checkpoints.append(step_exec['checkpoint_id'])

            self._executions[execution.execution_id] = execution

    def create_execution(
        self,
        workflow_name: str,
        inputs: Dict[str, Any]
    ) -> Execution:
        """Create a new execution."""
        execution = Execution(workflow_name, inputs)
        self._executions[execution.execution_id] = execution

        # Save to database immediately
        execution._save_to_db(self.db)

        return execution

    def get_execution(self, execution_id: str) -> Execution | None:
        """Get an execution by ID."""
        return self._executions.get(execution_id)

    def get_executions_for_workflow(self, workflow_name: str) -> List[Execution]:
        """Get all executions for a workflow."""
        return [
            ex for ex in self._executions.values()
            if ex.workflow_name == workflow_name
        ]

    def get_all_executions(self) -> List[Execution]:
        """Get all executions."""
        return list(self._executions.values())

    def set_current_execution(self, execution_id: str) -> None:
        """Set the current execution context."""
        self._current_execution_id = execution_id

    def get_current_execution(self) -> Execution | None:
        """Get the current execution."""
        if self._current_execution_id:
            return self._executions.get(self._current_execution_id)
        return None

    def clear_current_execution(self) -> None:
        """Clear the current execution context."""
        self._current_execution_id = None

    def delete_execution(self, execution_id: str) -> bool:
        """Delete an execution from memory and database.

        Args:
            execution_id: Execution ID to delete

        Returns:
            True if execution was deleted, False if not found
        """
        # Delete from database
        deleted = self.db.delete_execution(execution_id)

        # Remove from memory if present
        if execution_id in self._executions:
            del self._executions[execution_id]

        # Clear current execution if it's the one being deleted
        if self._current_execution_id == execution_id:
            self._current_execution_id = None

        return deleted


# Global execution manager instance
_execution_manager: ExecutionManager | None = None


def get_execution_manager() -> ExecutionManager:
    """Get the global execution manager (lazy singleton)."""
    global _execution_manager
    if _execution_manager is None:
        _execution_manager = ExecutionManager()
    return _execution_manager
