"""Global decorators and execution context for workflow steps."""

import functools
import asyncio
from typing import Any, Callable, ParamSpec, TypeVar, Coroutine, overload, TypeAlias


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

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


def set_execution_context(execution_id: str, workflow_name: str) -> None:
    """Set the current execution context.

    Args:
        execution_id: ID of the current execution
        workflow_name: Name of the workflow being executed (unused but kept for compatibility)
    """
    from ombra.execution import get_execution_manager
    execution_manager = get_execution_manager()
    execution_manager.set_current_execution(execution_id)


def clear_execution_context() -> None:
    """Clear the execution context."""
    from ombra.execution import get_execution_manager
    execution_manager = get_execution_manager()
    execution_manager.clear_current_execution()


def get_execution_context() -> tuple[str | None, str | None]:
    """Get the current execution context.

    Returns:
        Tuple of (execution_id, workflow_name)
    """
    from ombra.execution import get_execution_manager
    execution_manager = get_execution_manager()
    execution = execution_manager.get_current_execution()

    if execution:
        return execution.execution_id, execution.workflow_name

    return None, None


def step(name: str | None = None) -> Callable[[AsyncFunction[P, R]], AsyncFunction[P, R]]:
    """Global decorator to mark a function as a workflow step.

    When a step executes:
    - Captures inputs and outputs
    - Saves checkpoint if execution context is active
    - Tracks step execution order

    Args:
        name: Optional custom name for the step. Defaults to function name.

    Usage:
        @step()
        async def my_function(x: int) -> int:
            return x * 2
    """
    def decorator(func: AsyncFunction[P, R]) -> AsyncFunction[P, R]:
        step_name = name or func.__name__

        # Enforce async - step functions MUST be async
        if not asyncio.iscoroutinefunction(func):
            raise TypeError(
                f"Step function '{step_name}' must be async. "
                f"Use 'async def {func.__name__}' instead of 'def {func.__name__}'"
            )

        # Store metadata on the function
        func.__step_name__ = step_name
        func.__is_step__ = True

        # Create async wrapper
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            execution_id, workflow_name = get_execution_context()
            inputs = {"args": args, "kwargs": kwargs}

            # Check replay mode
            if execution_id and workflow_name:
                from ombra.execution import get_execution_manager
                execution_manager = get_execution_manager()
                execution = execution_manager.get_execution(execution_id)

                if execution and execution.replay_mode:
                    if step_name == execution.replay_from_step:
                        execution.replay_mode = False
                        print(f"🎯 Reached target step '{step_name}' - resuming execution")
                    else:
                        if step_name in execution.replay_checkpoint_cache:
                            cached_result = execution.replay_checkpoint_cache[step_name]
                            print(f"⚡ Skipped '{step_name}' (using cached result)")
                            execution.add_step_execution(step_name, inputs, cached_result)
                            return cached_result

            # Execute async function
            error = None
            result = None
            try:
                result = await func(*args, **kwargs)
            except Exception as e:
                error = str(e)

            # Save checkpoint
            _save_checkpoint(execution_id, workflow_name, step_name, inputs, result, error, async_wrapper, func)

            if error:
                raise Exception(error)
            return result

        return async_wrapper

    return decorator


def _save_checkpoint(execution_id, workflow_name, step_name, inputs, result, error, wrapper, func):
    """Save checkpoint for a step execution."""
    if execution_id and workflow_name:
        from ombra.execution import get_execution_manager

        execution_manager = get_execution_manager()
        execution = execution_manager.get_execution(execution_id)

        if execution:
            import ombra
            workflow = ombra.get_workflow(workflow_name)

            if workflow:
                if step_name not in workflow._steps:
                    workflow._steps[step_name] = wrapper

                checkpoint_id = workflow.checkpoint_manager.save(
                    step_name=step_name,
                    inputs=inputs,
                    outputs=result,
                    metadata={
                        "execution_id": execution_id,
                        "func_name": func.__name__
                    }
                )

                execution.add_checkpoint(checkpoint_id)
                execution.add_step_execution(step_name, inputs, result, error=error)

                exec_manager = get_execution_manager()
                sequence_order = len(execution.step_executions) - 1
                exec_manager.db.save_step_execution(
                    execution_id=execution_id,
                    step_name=step_name,
                    sequence_order=sequence_order,
                    checkpoint_id=checkpoint_id
                )

                if error:
                    print(f"✗ Step '{step_name}' failed → checkpoint: {checkpoint_id}")
                else:
                    print(f"✓ Step '{step_name}' completed → checkpoint: {checkpoint_id}")
