"""Core workflow engine with step decorator and time travel."""

import functools
from typing import Callable, Any, ParamSpec, TypeVar, Coroutine, TypeAlias
from .checkpoint import CheckpointManager


P = ParamSpec('P')
R = TypeVar('R')

# Type alias for async functions - makes error messages clearer
AsyncFunction: TypeAlias = Callable[P, Coroutine[Any, Any, R]]


class Workflow:
    """Main workflow orchestrator with checkpointing capabilities."""

    def __init__(self, name: str, db_path: str = ".ombra_executions.db", description: str = ""):
        self.name = name
        self.description = description
        self.checkpoint_manager = CheckpointManager(name, db_path)
        self._steps: dict[str, Callable] = {}
        self._runner: Callable | None = None

        # Auto-register this workflow
        import ombra
        ombra.register_workflow(self)

    def __call__(self, *args, **kwargs):
        """Make the workflow callable - executes the runner function.

        This allows workflows to be called directly:
            result = my_workflow(x=1, y=2)
        """
        if not self._runner:
            raise ValueError(f"No runner function defined for workflow '{self.name}'")
        return self._runner(*args, **kwargs)

    def runner(self):
        """Decorator to mark a function as the workflow runner (entry point).

        The runner function defines the complete workflow execution logic.
        Its parameters become the API endpoint parameters.

        Usage:
            @workflow.runner()
            def run(query: str) -> dict:
                data = fetch_data(query)
                result = process_data(data)
                return result
        """
        def decorator(func: Callable) -> Callable:
            self._runner = func
            return func
        return decorator

    def list_checkpoints(self, step_name: str | None = None):
        """List all checkpoints for this workflow.

        Args:
            step_name: Optional filter to show only checkpoints for a specific step.
        """
        checkpoints = self.checkpoint_manager.list_checkpoints(step_name)

        if not checkpoints:
            print(f"No checkpoints found for workflow '{self.name}'")
            return

        print(f"\n📦 Checkpoints for workflow '{self.name}':")
        print("-" * 60)

        for cp in checkpoints:
            print(f"  {cp['checkpoint_id']}")
            print(f"    Step: {cp['step_name']}")
            print(f"    Time: {cp['timestamp']}")
            print(f"    Inputs: {cp['inputs']}")
            print(f"    Outputs: {cp['outputs']}")
            print()

    def load_checkpoint(self, checkpoint_id: str) -> dict:
        """Load a specific checkpoint by ID.

        Args:
            checkpoint_id: The ID of the checkpoint to load.

        Returns:
            Dictionary with checkpoint data (inputs, outputs, metadata).
        """
        checkpoint = self.checkpoint_manager.load(checkpoint_id)
        print(f"📂 Loaded checkpoint: {checkpoint_id}")
        return checkpoint

    def get_last_input(self, step_name: str) -> Any:
        """Get the inputs from the most recent execution of a step.

        Args:
            step_name: Name of the step to get inputs for.

        Returns:
            The input arguments that were passed to the step.
        """
        latest = self.checkpoint_manager.get_latest(step_name)

        if not latest:
            raise ValueError(f"No checkpoints found for step '{step_name}'")

        inputs = latest["inputs"]
        print(f"🔄 Retrieved last inputs for '{step_name}': {inputs}")
        return inputs

    def resume_from(self, step_name: str, modify_inputs: dict | None = None):
        """Resume execution from a specific step with optional input modification.

        Args:
            step_name: Name of the step to resume from.
            modify_inputs: Optional dictionary to override inputs (kwargs only).

        Returns:
            Result of executing the step with the loaded (or modified) inputs.
        """
        if step_name not in self._steps:
            raise ValueError(f"Step '{step_name}' not found in workflow")

        # Get the latest inputs for this step
        inputs = self.get_last_input(step_name)
        args = inputs["args"]
        kwargs = inputs["kwargs"]

        # Apply modifications if provided
        if modify_inputs:
            # When modifying, use only kwargs to avoid conflicts
            kwargs.update(modify_inputs)
            args = ()  # Clear positional args to use kwargs only
            print(f"🔧 Modified inputs: {modify_inputs}")

        # Execute the step
        func = self._steps[step_name]
        return func(*args, **kwargs)

    async def resume_workflow_from(self, step_name: str, runner_inputs: dict):
        """Resume ENTIRE workflow from a specific step to end.

        This loads checkpoints for all prior steps and re-executes the runner,
        skipping steps until it reaches step_name, then continues normally.

        Args:
            step_name: Name of the step to resume from
            runner_inputs: Original inputs to the runner function

        Returns:
            Result of executing the workflow from the target step to end
        """
        if not self._runner:
            raise ValueError("No runner function defined for this workflow")

        # Note: We don't check if step_name is in self._steps here because
        # after a code reload, steps haven't been registered yet.
        # They will auto-register during execution when the runner calls them.

        # Get execution manager
        from ombra.execution import get_execution_manager
        execution_manager = get_execution_manager()

        # Load checkpoints from the most recent execution
        executions = execution_manager.get_executions_for_workflow(self.name)
        if not executions:
            raise ValueError("No previous executions found to resume from")

        # Find the most recent completed execution with steps
        completed_executions = [
            ex for ex in executions
            if ex.step_executions and ex.status.value in ["completed", "failed"]
        ]
        if not completed_executions:
            raise ValueError("No completed executions found with step history")

        last_execution = sorted(
            completed_executions,
            key=lambda ex: ex.started_at,
            reverse=True
        )[0]

        # Build checkpoint cache: load outputs for steps BEFORE the target
        checkpoint_cache = {}
        reached_target = False

        for step_exec in last_execution.step_executions:
            current_step = step_exec['step_name']

            if current_step == step_name:
                reached_target = True
                break  # Don't cache the target step or anything after

            # Cache this step's output
            checkpoint_cache[current_step] = step_exec['outputs']
            print(f"📦 Cached result for '{current_step}'")

        if not reached_target:
            raise ValueError(f"Step '{step_name}' not found in last execution's step history")

        # Create a new execution in replay mode
        execution = execution_manager.create_execution(
            workflow_name=self.name,
            inputs={
                "resumed_from_step": step_name,
                "original_runner_inputs": runner_inputs
            }
        )
        execution.replay_mode = True
        execution.replay_from_step = step_name
        execution.replay_checkpoint_cache = checkpoint_cache

        # Start execution and set context
        execution.start(db=execution_manager.db)
        from ombra.decorators import set_execution_context, clear_execution_context

        try:
            set_execution_context(execution.execution_id, self.name)

            print(f"🔄 Re-executing workflow from '{step_name}' to end...")

            # Execute the runner - it will skip cached steps and resume from target
            result = await self._runner(**runner_inputs)

            execution.complete(result, db=execution_manager.db)
            print(f"✅ Workflow completed from '{step_name}' to end")

            return result

        except Exception as e:
            execution.fail(str(e), db=execution_manager.db)
            raise

        finally:
            clear_execution_context()


def _extract_fastapi_annotation(value: Any) -> tuple[str | None, dict]:
    """Extract FastAPI parameter annotation information.

    Detects FastAPI parameter types like File, Form, Body, Query, Path, Header, Cookie.

    Args:
        value: The default value to analyze

    Returns:
        Tuple of (annotation_type, metadata_dict)
        e.g., ("File", {...}) or (None, {})
    """
    value_type = type(value).__name__
    value_module = type(value).__module__ if hasattr(value, '__module__') else None

    # Check if it's from fastapi.params
    if value_module and 'fastapi' in value_module:
        # Common FastAPI parameter types
        fastapi_types = ['File', 'Form', 'Body', 'Query', 'Path', 'Header', 'Cookie', 'Depends']

        if value_type in fastapi_types:
            # Extract any attributes that might be useful
            metadata = {}
            if hasattr(value, 'default'):
                metadata['has_default'] = value.default is not ...
            return value_type, metadata

    return None, {}


def _make_json_serializable(value: Any) -> Any:
    """Convert a value to JSON-serializable format.

    Handles FastAPI special types like File, Form, Body, etc.

    Args:
        value: The value to make JSON serializable

    Returns:
        A JSON-serializable representation of the value
    """
    import json

    # Try to serialize directly
    try:
        json.dumps(value)
        return value
    except (TypeError, ValueError):
        pass

    # Handle FastAPI special types
    annotation_type, _ = _extract_fastapi_annotation(value)
    if annotation_type:
        return f"{annotation_type}(...)"

    # For other non-serializable types, return string representation
    return str(value)


def workflow(name: str, description: str = "", db_path: str = ".ombra_executions.db") -> Callable[[AsyncFunction[P, R]], Workflow]:
    """Decorator to define a workflow with a single decorator.

    This is a simplified API that combines workflow creation and runner definition.

    Args:
        name: Workflow name
        description: Optional workflow description
        db_path: Path to SQLite database file

    Usage:
        @workflow(name="my_workflow", description="Does something cool")
        async def my_workflow(x: int, y: int) -> str:
            result = await some_step(x, y)
            return result

    The decorated function becomes the workflow runner and the workflow
    is returned as a callable object that can be invoked directly.
    """
    def decorator(func: AsyncFunction[P, R]) -> Workflow:
        import inspect
        import asyncio
        from typing import get_type_hints

        # Enforce async - workflow runner MUST be async
        if not asyncio.iscoroutinefunction(func):
            raise TypeError(
                f"Workflow runner '{name}' must be async. "
                f"Use 'async def {func.__name__}' instead of 'def {func.__name__}'"
            )

        # Create workflow instance
        wf = Workflow(name=name, description=description, db_path=db_path)

        # Set the decorated function as the runner
        wf._runner = func

        # Extract parameter signature and persist to database
        sig = inspect.signature(func)
        try:
            type_hints = get_type_hints(func)
        except:
            type_hints = {}

        parameters = []
        for param_name, param in sig.parameters.items():
            # Handle default value - make it JSON serializable and extract FastAPI annotation
            if param.default == inspect.Parameter.empty:
                default_value = None
                annotation_type = None
                annotation_metadata = {}
            else:
                default_value = _make_json_serializable(param.default)
                annotation_type, annotation_metadata = _extract_fastapi_annotation(param.default)

            param_info = {
                'name': param_name,
                'type': type_hints.get(param_name, 'Any').__name__ if param_name in type_hints else 'Any',
                'required': param.default == inspect.Parameter.empty,
                'default': default_value
            }

            # Add FastAPI annotation info if present
            if annotation_type:
                param_info['fastapi_annotation'] = annotation_type
                param_info['annotation_metadata'] = annotation_metadata

            parameters.append(param_info)

        # Save schema to database
        from ombra.execution import get_execution_manager
        exec_manager = get_execution_manager()
        exec_manager.db.save_workflow_schema(name, description, parameters)

        # Return the workflow (which is now callable via __call__)
        return wf

    return decorator
