"""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,
        step_input_overrides: Dict[str, Any] | None = None
    ):
        self.execution_id = execution_id or str(uuid.uuid4())
        self.workflow_name = workflow_name
        self.inputs = inputs
        self.step_input_overrides = step_input_overrides or {}
        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] = {}
        
        # Track call counts for unique step identification: parent_name -> { step_name -> count }
        self.step_call_counts: Dict[str, Dict[str, int]] = {}

    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, parent_checkpoint_id: str | None = None, checkpoint_id: str | None = None, parent_step_name: str | None = None, cache_key: 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)
            parent_checkpoint_id: ID of the parent step's checkpoint (if nested)
            checkpoint_id: ID of this step's checkpoint
            parent_step_name: Name of the parent step (optional, for cache key generation)
            cache_key: Pre-calculated cache key (optional). If provided, skips counter increment.
        """
        # If cache_key is not provided, calculate it (and increment counters)
        if not cache_key:
            # Calculate cache key for unique identification
            # We prefer the passed parent_step_name if available
            parent_name = parent_step_name or "ROOT"
            
            # Fallback logic (mostly for replay skipping where we might not pass it explicitly yet, though we should)
            if not parent_step_name and parent_checkpoint_id:
                # Find parent in existing steps
                for step in self.step_executions:
                    if step.get("checkpoint_id") == parent_checkpoint_id:
                        parent_name = step["step_name"]
                        break
            
            # Update call counts
            if parent_name not in self.step_call_counts:
                self.step_call_counts[parent_name] = {}
            
            current_count = self.step_call_counts[parent_name].get(step_name, 0)
            self.step_call_counts[parent_name][step_name] = current_count + 1
            
            cache_key = f"{parent_name}|{step_name}|{current_count}"

        self.step_executions.append({
            "step_name": step_name,
            "inputs": _make_serializable(inputs),
            "outputs": _make_serializable(outputs),
            "error": error,
            "timestamp": datetime.now().isoformat(),
            "parent_checkpoint_id": parent_checkpoint_id,
            "checkpoint_id": checkpoint_id,
            "cache_key": cache_key
        })

    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)
            
            # Helper map to resolve parent names for cache_key generation
            checkpoint_id_to_name = {}
            
            # First pass: build ID to Name map (because parents might be later in the list if they finish last)
            for step_exec in step_execs:
                if step_exec['checkpoint_id']:
                    checkpoint_id_to_name[step_exec['checkpoint_id']] = step_exec['step_name']

            # Second pass: build execution objects
            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:
                        
                        # Reconstruct cache_key
                        parent_checkpoint_id = step_exec.get('parent_checkpoint_id')
                        parent_name = "ROOT"
                        if parent_checkpoint_id and parent_checkpoint_id in checkpoint_id_to_name:
                            parent_name = checkpoint_id_to_name[parent_checkpoint_id]
                        
                        if parent_name not in execution.step_call_counts:
                            execution.step_call_counts[parent_name] = {}
                        
                        current_count = execution.step_call_counts[parent_name].get(step_exec['step_name'], 0)
                        execution.step_call_counts[parent_name][step_exec['step_name']] = current_count + 1
                        
                        cache_key = f"{parent_name}|{step_exec['step_name']}|{current_count}"

                        execution.step_executions.append({
                            'step_name': step_exec['step_name'],
                            'inputs': checkpoint['inputs'],
                            'outputs': checkpoint['outputs'],
                            'timestamp': step_exec['timestamp'],
                            'checkpoint_id': step_exec['checkpoint_id'],
                            'parent_checkpoint_id': step_exec.get('parent_checkpoint_id'),
                            'cache_key': cache_key
                        })
                        execution.checkpoints.append(step_exec['checkpoint_id'])

            self._executions[execution.execution_id] = execution

    def create_execution(
        self,
        workflow_name: str,
        inputs: Dict[str, Any],
        step_input_overrides: Dict[str, Any] | None = None
    ) -> Execution:
        """Create a new execution."""
        # Pre-process inputs to persist files
        from .storage import get_storage_manager
        storage = get_storage_manager()
        
        def process_value(value: Any) -> Any:
            # Check for RehydratedFile (has to_metadata)
            if hasattr(value, 'to_metadata') and callable(value.to_metadata):
                # It's already a rehydrated file (from "Run Again"), persist it as a new snapshot
                return storage.save_file(value)

            # Check for Starlette/FastAPI UploadFile
            # We use the same logic as _make_serializable to ensure consistency
            obj_type = type(value).__name__
            obj_module = type(value).__module__ if hasattr(value, '__module__') else None

            is_upload_file = False
            if obj_type == 'UploadFile' or (obj_module and 'starlette' in obj_module and 'UploadFile' in obj_type):
                is_upload_file = True
            
            # Fallback duck typing (has filename and file)
            if not is_upload_file and hasattr(value, 'filename') and hasattr(value, 'file') and not isinstance(value, (dict, list, str, int, float, bool, type(None))):
                is_upload_file = True

            if is_upload_file:
                return storage.save_file(value)
            
            if isinstance(value, dict):
                return {k: process_value(v) for k, v in value.items()}
            
            if isinstance(value, list):
                return [process_value(v) for v in value]
                
            return value

        processed_inputs = process_value(inputs)

        execution = Execution(workflow_name, processed_inputs, step_input_overrides=step_input_overrides)
        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
