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

import functools
import asyncio
import contextvars
from datetime import datetime
from typing import Any, Callable, ParamSpec, TypeVar, Coroutine, overload, TypeAlias, get_type_hints, get_origin, get_args
from ombra.logger import logger


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

# Type alias for async functions
AsyncFunction: TypeAlias = Callable[P, Coroutine[Any, Any, R]]

# Context variable to track the stack of active step execution IDs (checkpoints)
# Stores dicts: {"id": checkpoint_id, "name": step_name}
_active_checkpoint_stack: contextvars.ContextVar[list[dict[str, str]]] = contextvars.ContextVar("active_checkpoint_stack", default=[])


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()
            
            # Bind arguments to signature to get meaningful input names
            import inspect
            sig = inspect.signature(func)
            bound = sig.bind(*args, **kwargs)
            bound.apply_defaults()
            
            # Convert to dict and filter out 'self' if present
            inputs = dict(bound.arguments)
            if 'self' in inputs:
                del inputs['self']

            # Generate checkpoint ID upfront
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            checkpoint_id = f"{step_name}_{timestamp}"

            # Manage nesting stack
            stack = _active_checkpoint_stack.get()
            parent_entry = stack[-1] if stack else None
            parent_checkpoint_id = parent_entry["id"] if parent_entry else None
            parent_step_name = parent_entry["name"] if parent_entry else "ROOT"
            
            # Push current ID and Name to stack
            token = _active_checkpoint_stack.set(stack + [{"id": checkpoint_id, "name": step_name}])

            try:
                # Check replay mode and overrides
                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:
                        # Increment call count for this step within its parent
                        if parent_step_name not in execution.step_call_counts:
                            execution.step_call_counts[parent_step_name] = {}
                        
                        current_count = execution.step_call_counts[parent_step_name].get(step_name, 0)
                        execution.step_call_counts[parent_step_name][step_name] = current_count + 1
                        
                        # Generate unique cache key: Parent|Step|Index
                        cache_key = f"{parent_step_name}|{step_name}|{current_count}"

                        # 1. Bind arguments first to get a unified view
                        sig = inspect.signature(func)
                        bound = sig.bind(*args, **kwargs)
                        bound.apply_defaults()

                        # 2. Check for input overrides (Time Travel modification)
                        overrides = None
                        if execution.step_input_overrides:
                            if cache_key in execution.step_input_overrides:
                                overrides = execution.step_input_overrides[cache_key]
                                logger.info(f"🔧 Applying input overrides for '{step_name}' (Key: {cache_key}): {overrides}")
                            elif step_name in execution.step_input_overrides:
                                overrides = execution.step_input_overrides[step_name]
                                logger.info(f"🔧 Applying input overrides for '{step_name}': {overrides}")

                        if overrides:
                            # Apply overrides to bound arguments
                            bound.arguments.update(overrides)
                            
                            # Update our inputs tracking dict as well for the checkpoint
                            inputs.update(overrides)

                        # 3. Auto-convert dicts to Pydantic models (for BOTH regular args and overrides)
                        try:
                            type_hints = get_type_hints(func)
                            for key, value in bound.arguments.items():
                                if key not in type_hints:
                                    continue
                                    
                                param_type = type_hints[key]
                                origin = get_origin(param_type)
                                args_types = get_args(param_type)

                                # Case 1: Direct Pydantic Model (e.g. arg: MyModel)
                                if isinstance(value, dict):
                                    if hasattr(param_type, 'model_validate'):
                                        try:
                                            bound.arguments[key] = param_type.model_validate(value)
                                        except Exception as e:
                                            logger.warning(f"⚠️ Failed to convert arg '{key}' to Pydantic model: {e}")
                                    elif hasattr(param_type, 'parse_obj'):
                                        try:
                                            bound.arguments[key] = param_type.parse_obj(value)
                                        except Exception as e:
                                            logger.warning(f"⚠️ Failed to convert arg '{key}' to Pydantic model: {e}")
                                
                                # Case 2: List of Pydantic Models (e.g. arg: List[MyModel])
                                elif origin is list and args_types and isinstance(value, list):
                                    inner_type = args_types[0]
                                    if hasattr(inner_type, 'model_validate') or hasattr(inner_type, 'parse_obj'):
                                        new_list = []
                                        for item in value:
                                            if isinstance(item, dict):
                                                try:
                                                    if hasattr(inner_type, 'model_validate'):
                                                        new_list.append(inner_type.model_validate(item))
                                                    else:
                                                        new_list.append(inner_type.parse_obj(item))
                                                except Exception as e:
                                                    logger.warning(f"⚠️ Failed to convert list item in '{key}': {e}")
                                                    new_list.append(item)
                                            else:
                                                new_list.append(item)
                                        bound.arguments[key] = new_list

                        except Exception as e:
                            # Don't fail execution if type checking fails, just warn
                            # logger.warning(f"⚠️ Type hint inspection failed: {e}")
                            pass

                        # 4. Update args/kwargs for the actual call
                        args = bound.args
                        kwargs = bound.kwargs

                        # Handle Replay
                        if execution.replay_mode:
                            # Check cache_key first (more specific), then step_name
                            stop_replay = False
                            if execution.replay_from_step == cache_key:
                                stop_replay = True
                            elif execution.replay_from_step == step_name:
                                stop_replay = True
                            
                            if stop_replay:
                                execution.replay_mode = False
                                logger.info(f"🎯 Reached target step '{step_name}' (Key: {cache_key}) - resuming execution")
                            else:
                                if cache_key in execution.replay_checkpoint_cache:
                                    cached_result = execution.replay_checkpoint_cache[cache_key]
                                    logger.info(f"⚡ Skipped '{step_name}' (using cached result from '{cache_key}')")
                                    
                                    # Record skipped step with correct parent linkage
                                    execution.add_step_execution(
                                        step_name, 
                                        inputs, 
                                        cached_result, 
                                        parent_checkpoint_id=parent_checkpoint_id,
                                        checkpoint_id=checkpoint_id 
                                    )
                                    return cached_result

                # Execute async function
                error = None
                result = None
                try:
                    result = await func(*args, **kwargs)
                except Exception as e:
                    import traceback
                    error = traceback.format_exc()
                    
                    _save_checkpoint(execution_id, workflow_name, step_name, inputs, result, error, async_wrapper, func, checkpoint_id, parent_checkpoint_id, parent_step_name, cache_key)

                    if error:
                        raise e
                
                # Save checkpoint (success case)
                _save_checkpoint(execution_id, workflow_name, step_name, inputs, result, error, async_wrapper, func, checkpoint_id, parent_checkpoint_id, parent_step_name, cache_key)

                return result
            finally:
                # Always pop stack
                _active_checkpoint_stack.reset(token)

        return async_wrapper

    return decorator


def _save_checkpoint(execution_id, workflow_name, step_name, inputs, result, error, wrapper, func, checkpoint_id, parent_checkpoint_id, parent_step_name, cache_key):
    """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

                # Save using pre-generated ID
                saved_id = workflow.checkpoint_manager.save(
                    step_name=step_name,
                    inputs=inputs,
                    outputs=result,
                    metadata={
                        "execution_id": execution_id,
                        "func_name": func.__name__
                    },
                    checkpoint_id=checkpoint_id
                )

                execution.add_checkpoint(saved_id)
                
                # Pass parent_checkpoint_id, checkpoint_id, parent_step_name AND cache_key
                execution.add_step_execution(
                    step_name, 
                    inputs, 
                    result, 
                    error=error, 
                    parent_checkpoint_id=parent_checkpoint_id, 
                    checkpoint_id=saved_id,
                    parent_step_name=parent_step_name,
                    cache_key=cache_key
                )

                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=saved_id,
                    parent_checkpoint_id=parent_checkpoint_id
                )

                if error:
                    logger.error(f"✗ Step '{step_name}' failed → checkpoint: {saved_id}")
                else:
                    logger.info(f"✓ Step '{step_name}' completed → checkpoint: {saved_id}")
