"""
Workflow Replay System

Enables re-execution of workflows from specific steps with modified context.
Supports debugging, testing, and recovery scenarios.

This is a capability n8n lacks - step-level replay with context modification.

Design principles:
- Non-destructive: Original evidence preserved
- Flexible: Replay from any step with any context modification
- Traceable: Replay executions are linked to originals
"""

import copy
import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set

logger = logging.getLogger(__name__)


class ReplayMode(str, Enum):
    """Replay execution modes"""
    FROM_STEP = "from_step"          # Continue from step to end
    SINGLE_STEP = "single_step"      # Execute only the specified step
    RANGE = "range"                  # Execute step range
    MODIFIED = "modified"            # Re-run only modified steps


@dataclass
class ReplayConfig:
    """
    Configuration for replay execution.

    Attributes:
        mode: Replay mode
        start_step_id: Step to start from
        end_step_id: Optional step to end at
        modified_context: Context modifications to apply
        skip_steps: Steps to skip during replay
        breakpoints: Steps to pause at
        dry_run: If True, validate only without executing
    """
    mode: ReplayMode = ReplayMode.FROM_STEP
    start_step_id: str = ""
    end_step_id: Optional[str] = None
    modified_context: Dict[str, Any] = field(default_factory=dict)
    skip_steps: Set[str] = field(default_factory=set)
    breakpoints: Set[str] = field(default_factory=set)
    dry_run: bool = False


@dataclass
class ReplayResult:
    """
    Result of a replay execution.

    Attributes:
        ok: Whether replay completed successfully
        execution_id: New execution ID for this replay
        original_execution_id: Original execution being replayed
        start_step: Step replay started from
        end_step: Step replay ended at
        steps_executed: Number of steps executed
        steps_skipped: Number of steps skipped
        context: Final context after replay
        error: Error message if failed
        duration_ms: Total replay duration
    """
    ok: bool
    execution_id: str
    original_execution_id: str
    start_step: str
    end_step: Optional[str] = None
    steps_executed: int = 0
    steps_skipped: int = 0
    context: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None
    duration_ms: int = 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "ok": self.ok,
            "execution_id": self.execution_id,
            "original_execution_id": self.original_execution_id,
            "start_step": self.start_step,
            "end_step": self.end_step,
            "steps_executed": self.steps_executed,
            "steps_skipped": self.steps_skipped,
            "context": self.context,
            "error": self.error,
            "duration_ms": self.duration_ms,
        }


class ReplayManager:
    """
    Manages workflow replay operations.

    Loads evidence from previous executions and enables re-execution
    from any point with optional context modifications.

    Usage:
        manager = ReplayManager(evidence_store)

        # Replay from specific step
        result = await manager.replay_from_step(
            execution_id="exec_123",
            step_id="step_5",
            modified_context={"user_id": "new_user"},
            workflow_executor=executor_func,
        )

        # Validate replay without executing
        validation = await manager.validate_replay(
            execution_id="exec_123",
            step_id="step_5",
        )
    """

    def __init__(
        self,
        evidence_path: Path,
        max_replay_history: int = 100,
    ):
        """
        Initialize replay manager.

        Args:
            evidence_path: Path to evidence storage
            max_replay_history: Maximum replay history to maintain
        """
        self.evidence_path = Path(evidence_path)
        self.max_replay_history = max_replay_history
        self._replay_history: List[ReplayResult] = []

    async def load_execution_state(
        self,
        execution_id: str,
        step_id: str,
    ) -> Optional[Dict[str, Any]]:
        """
        Load execution state at a specific step.

        Args:
            execution_id: Original execution ID
            step_id: Step to load state at

        Returns:
            Context state at that step, or None if not found
        """
        jsonl_path = self.evidence_path / execution_id / "evidence.jsonl"

        if not jsonl_path.exists():
            logger.warning(f"Evidence not found for {execution_id}")
            return None

        try:
            with open(jsonl_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    data = json.loads(line)
                    if data.get('step_id') == step_id:
                        return {
                            "context_before": data.get('context_before', {}),
                            "context_after": data.get('context_after', {}),
                            "step_index": data.get('step_index'),
                            "module_id": data.get('module_id'),
                            "status": data.get('status'),
                            "output": data.get('output', {}),
                        }

        except Exception as e:
            logger.error(f"Failed to load state for {execution_id}/{step_id}: {e}")

        return None

    async def load_execution_steps(
        self,
        execution_id: str,
    ) -> List[Dict[str, Any]]:
        """
        Load all steps from an execution.

        Args:
            execution_id: Execution ID

        Returns:
            List of step data in execution order
        """
        jsonl_path = self.evidence_path / execution_id / "evidence.jsonl"

        if not jsonl_path.exists():
            return []

        steps = []
        try:
            with open(jsonl_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line:
                        steps.append(json.loads(line))
        except Exception as e:
            logger.error(f"Failed to load steps for {execution_id}: {e}")

        return steps

    async def validate_replay(
        self,
        execution_id: str,
        step_id: str,
        config: Optional[ReplayConfig] = None,
    ) -> Dict[str, Any]:
        """
        Validate that a replay is possible.

        Args:
            execution_id: Execution to replay from
            step_id: Step to start from
            config: Optional replay configuration

        Returns:
            Validation result with issues if any
        """
        config = config or ReplayConfig(start_step_id=step_id)
        issues = []
        warnings = []

        # Check evidence exists
        exec_dir = self.evidence_path / execution_id
        if not exec_dir.exists():
            issues.append(f"Execution {execution_id} not found")
            return {
                "valid": False,
                "issues": issues,
                "warnings": warnings,
            }

        # Load steps
        steps = await self.load_execution_steps(execution_id)
        if not steps:
            issues.append("No step evidence found")
            return {
                "valid": False,
                "issues": issues,
                "warnings": warnings,
            }

        # Check start step exists
        step_ids = [s.get('step_id') for s in steps]
        if step_id not in step_ids:
            issues.append(f"Step {step_id} not found in execution")
            return {
                "valid": False,
                "issues": issues,
                "warnings": warnings,
            }

        # Check end step if specified
        if config.end_step_id and config.end_step_id not in step_ids:
            issues.append(f"End step {config.end_step_id} not found")

        # Check skip steps exist
        for skip_id in config.skip_steps:
            if skip_id not in step_ids:
                warnings.append(f"Skip step {skip_id} not found")

        # Load state at start step
        state = await self.load_execution_state(execution_id, step_id)
        if not state:
            issues.append(f"Could not load state at step {step_id}")

        # Validate modified context keys
        if state and config.modified_context:
            original_keys = set(state.get('context_before', {}).keys())
            for key in config.modified_context:
                if key not in original_keys:
                    warnings.append(f"Modified key '{key}' not in original context")

        return {
            "valid": len(issues) == 0,
            "issues": issues,
            "warnings": warnings,
            "step_count": len(steps),
            "start_step_index": step_ids.index(step_id) if step_id in step_ids else None,
            "available_steps": step_ids,
        }

    async def prepare_replay_context(
        self,
        execution_id: str,
        step_id: str,
        modified_context: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Prepare context for replay execution.

        Loads the context_before from the target step and applies
        any modifications.

        Args:
            execution_id: Original execution ID
            step_id: Step to start from
            modified_context: Context modifications to apply

        Returns:
            Prepared context for replay
        """
        state = await self.load_execution_state(execution_id, step_id)

        if not state:
            raise ValueError(f"Cannot load state for {execution_id}/{step_id}")

        # Start with context before the step
        context = copy.deepcopy(state.get('context_before', {}))

        # Apply modifications
        if modified_context:
            context.update(modified_context)

        return context

    async def replay_from_step(
        self,
        execution_id: str,
        step_id: str,
        workflow_executor: Callable,
        config: Optional[ReplayConfig] = None,
    ) -> ReplayResult:
        """
        Execute replay from a specific step.

        Args:
            execution_id: Original execution ID
            step_id: Step to start from
            workflow_executor: Async function to execute workflow
                Signature: async def executor(workflow, context, start_step, end_step) -> result
            config: Replay configuration

        Returns:
            ReplayResult with execution outcome
        """
        config = config or ReplayConfig(start_step_id=step_id)
        start_time = datetime.now()
        new_execution_id = f"replay_{uuid.uuid4().hex[:12]}"

        # Validate first
        validation = await self.validate_replay(execution_id, step_id, config)
        if not validation.get('valid'):
            return ReplayResult(
                ok=False,
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                error=f"Validation failed: {validation.get('issues')}",
            )

        if config.dry_run:
            return ReplayResult(
                ok=True,
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                context={"dry_run": True, "validation": validation},
            )

        try:
            # Prepare context
            context = await self.prepare_replay_context(
                execution_id,
                step_id,
                config.modified_context,
            )

            # Get workflow definition (stored in evidence metadata)
            workflow = await self._load_workflow_definition(execution_id)
            if not workflow:
                return ReplayResult(
                    ok=False,
                    execution_id=new_execution_id,
                    original_execution_id=execution_id,
                    start_step=step_id,
                    error="Workflow definition not found in evidence",
                )

            # Execute workflow from step
            result = await workflow_executor(
                workflow=workflow,
                context=context,
                start_step=step_id,
                end_step=config.end_step_id,
                skip_steps=config.skip_steps,
                breakpoints=config.breakpoints,
            )

            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)

            replay_result = ReplayResult(
                ok=result.get('ok', False),
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                end_step=config.end_step_id,
                steps_executed=result.get('steps_executed', 0),
                steps_skipped=len(config.skip_steps),
                context=result.get('context', {}),
                error=result.get('error'),
                duration_ms=duration_ms,
            )

            # Record in history
            self._add_to_history(replay_result)

            return replay_result

        except Exception as e:
            logger.error(f"Replay failed: {e}")
            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)

            return ReplayResult(
                ok=False,
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                error=str(e),
                duration_ms=duration_ms,
            )

    async def replay_single_step(
        self,
        execution_id: str,
        step_id: str,
        step_executor: Callable,
        modified_params: Optional[Dict[str, Any]] = None,
        modified_context: Optional[Dict[str, Any]] = None,
    ) -> ReplayResult:
        """
        Re-execute a single step.

        Useful for testing individual step behavior with different inputs.

        Args:
            execution_id: Original execution ID
            step_id: Step to re-execute
            step_executor: Async function to execute single step
                Signature: async def executor(module_id, params, context) -> result
            modified_params: Parameter modifications
            modified_context: Context modifications

        Returns:
            ReplayResult with step outcome
        """
        start_time = datetime.now()
        new_execution_id = f"replay_step_{uuid.uuid4().hex[:12]}"

        try:
            # Load step state
            state = await self.load_execution_state(execution_id, step_id)
            if not state:
                return ReplayResult(
                    ok=False,
                    execution_id=new_execution_id,
                    original_execution_id=execution_id,
                    start_step=step_id,
                    error=f"Step {step_id} not found",
                )

            # Prepare context
            context = copy.deepcopy(state.get('context_before', {}))
            if modified_context:
                context.update(modified_context)

            # Get step configuration from workflow
            workflow = await self._load_workflow_definition(execution_id)
            step_config = self._find_step_config(workflow, step_id)

            if not step_config:
                return ReplayResult(
                    ok=False,
                    execution_id=new_execution_id,
                    original_execution_id=execution_id,
                    start_step=step_id,
                    error=f"Step configuration not found for {step_id}",
                )

            # Merge parameters
            params = copy.deepcopy(step_config.get('params', {}))
            if modified_params:
                params.update(modified_params)

            # Execute step
            result = await step_executor(
                module_id=step_config.get('module'),
                params=params,
                context=context,
            )

            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)

            return ReplayResult(
                ok=result.get('ok', True),
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                end_step=step_id,
                steps_executed=1,
                context=result.get('context', context),
                error=result.get('error'),
                duration_ms=duration_ms,
            )

        except Exception as e:
            logger.error(f"Single step replay failed: {e}")
            return ReplayResult(
                ok=False,
                execution_id=new_execution_id,
                original_execution_id=execution_id,
                start_step=step_id,
                error=str(e),
            )

    async def compare_replay(
        self,
        original_execution_id: str,
        replay_execution_id: str,
    ) -> Dict[str, Any]:
        """
        Compare original and replay execution results.

        Args:
            original_execution_id: Original execution ID
            replay_execution_id: Replay execution ID

        Returns:
            Comparison results
        """
        original_steps = await self.load_execution_steps(original_execution_id)
        replay_steps = await self.load_execution_steps(replay_execution_id)

        original_by_id = {s['step_id']: s for s in original_steps}
        replay_by_id = {s['step_id']: s for s in replay_steps}

        differences = []
        for step_id in set(original_by_id.keys()) | set(replay_by_id.keys()):
            orig = original_by_id.get(step_id)
            repl = replay_by_id.get(step_id)

            if not orig:
                differences.append({
                    "step_id": step_id,
                    "type": "added_in_replay",
                })
            elif not repl:
                differences.append({
                    "step_id": step_id,
                    "type": "not_in_replay",
                })
            else:
                # Compare outcomes
                if orig.get('status') != repl.get('status'):
                    differences.append({
                        "step_id": step_id,
                        "type": "status_changed",
                        "original": orig.get('status'),
                        "replay": repl.get('status'),
                    })

                if orig.get('output') != repl.get('output'):
                    differences.append({
                        "step_id": step_id,
                        "type": "output_changed",
                        "original_keys": list(orig.get('output', {}).keys()),
                        "replay_keys": list(repl.get('output', {}).keys()),
                    })

        return {
            "original_execution_id": original_execution_id,
            "replay_execution_id": replay_execution_id,
            "original_step_count": len(original_steps),
            "replay_step_count": len(replay_steps),
            "differences": differences,
            "identical": len(differences) == 0,
        }

    async def _load_workflow_definition(
        self,
        execution_id: str,
    ) -> Optional[Dict[str, Any]]:
        """Load workflow definition from execution metadata"""
        meta_path = self.evidence_path / execution_id / "workflow.json"

        if meta_path.exists():
            try:
                with open(meta_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                logger.warning(f"Failed to load workflow definition: {e}")

        return None

    def _find_step_config(
        self,
        workflow: Optional[Dict[str, Any]],
        step_id: str,
    ) -> Optional[Dict[str, Any]]:
        """Find step configuration by ID"""
        if not workflow:
            return None

        for step in workflow.get('steps', []):
            if step.get('id') == step_id:
                return step

        return None

    def _add_to_history(self, result: ReplayResult) -> None:
        """Add replay result to history"""
        self._replay_history.append(result)

        # Trim history if needed
        if len(self._replay_history) > self.max_replay_history:
            self._replay_history = self._replay_history[-self.max_replay_history:]

    def get_replay_history(
        self,
        execution_id: Optional[str] = None,
    ) -> List[ReplayResult]:
        """Get replay history, optionally filtered by original execution"""
        if execution_id:
            return [
                r for r in self._replay_history
                if r.original_execution_id == execution_id
            ]
        return list(self._replay_history)


# =============================================================================
# Factory Functions
# =============================================================================

def create_replay_manager(
    evidence_path: Optional[Path] = None,
) -> ReplayManager:
    """
    Create a replay manager.

    Args:
        evidence_path: Path to evidence storage

    Returns:
        Configured ReplayManager
    """
    if evidence_path is None:
        evidence_path = Path("./evidence")
    return ReplayManager(evidence_path=evidence_path)
