"""Pipeline decorators for marking async functions as pipelines."""

import functools
import inspect
from collections.abc import Callable, Coroutine
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, overload

from alloy_runtime_sdk.pipeline.context import PipelineContext

if TYPE_CHECKING:
    from pydantic import BaseModel

P = ParamSpec("P")
T = TypeVar("T")
InputT = TypeVar("InputT", bound="BaseModel")

# Type alias for pipeline functions
PipelineFunc = Callable[[PipelineContext], Coroutine[Any, Any, T]]

# Registry of all pipeline functions discovered in a module
_pipeline_registry: dict[str, Callable[..., Coroutine[Any, Any, Any]]] = {}


def _is_pydantic_model(obj: Any) -> bool:
    """Check if obj is a Pydantic BaseModel class (not instance).

    Uses attribute detection rather than issubclass to handle module isolation
    issues when code is executed via exec() in sandboxed environments (Modal).
    """
    if not isinstance(obj, type):
        return False

    # Check for Pydantic v2 attributes (model_fields, model_json_schema)
    if hasattr(obj, "model_fields") and hasattr(obj, "model_json_schema"):
        return True

    # Check for Pydantic v1 attributes (__fields__, schema)
    if hasattr(obj, "__fields__") and hasattr(obj, "schema"):
        return True

    # Fallback to issubclass check for edge cases
    try:
        from pydantic import BaseModel

        return issubclass(obj, BaseModel)
    except (ImportError, TypeError):
        return False


# Overload 1: Direct decoration without parentheses - @pipeline
@overload
def pipeline(
    func: PipelineFunc[T],
) -> PipelineFunc[T]: ...


# Overload 2: With schema class - @pipeline(MySchema)
@overload
def pipeline(
    schema: type[InputT],
    *,
    name: str | None = None,
) -> Callable[[PipelineFunc[T]], PipelineFunc[T]]: ...


# Overload 3: With keyword arguments only - @pipeline(name="custom")
@overload
def pipeline(
    schema: None = None,
    *,
    name: str | None = None,
) -> Callable[[PipelineFunc[T]], PipelineFunc[T]]: ...


def pipeline(  # pyright: ignore[reportInconsistentOverload]
    func_or_schema: PipelineFunc[T] | type[InputT] | None = None,
    *,
    name: str | None = None,
) -> PipelineFunc[T] | Callable[[PipelineFunc[T]], PipelineFunc[T]]:
    """Decorator to mark an async function as a pipeline.

    Supports multiple calling patterns:

    1. Without arguments (no parentheses):
        @pipeline
        async def my_pipeline(ctx: PipelineContext):
            return {"result": "ok"}

    2. With a Pydantic schema class for input validation:
        @pipeline(MyConfig)
        async def my_pipeline(ctx: PipelineContext):
            config: MyConfig = ctx.config
            return {"topic": config.topic}

    3. With keyword arguments:
        @pipeline(name="custom-name")
        async def my_pipeline(ctx: PipelineContext):
            return {"result": "ok"}

    4. With schema and keyword arguments:
        @pipeline(MyConfig, name="custom-name")
        async def my_pipeline(ctx: PipelineContext):
            return {"topic": ctx.config.topic}

    The decorated function must:
    1. Be an async function (coroutine)
    2. Accept a PipelineContext as its first parameter
    3. Return the pipeline result

    Args:
        func_or_schema: Either the decorated function (when used without parentheses)
                       or a Pydantic BaseModel class for input validation.
        name: Optional custom name for the pipeline. Defaults to function name.

    Returns:
        The decorated function with pipeline metadata attached.

    Raises:
        TypeError: If the function is not async or has wrong signature.
    """
    # Determine what was passed as the first argument
    schema_class: type[InputT] | None = None
    direct_func: PipelineFunc[T] | None = None

    if func_or_schema is not None:
        if callable(func_or_schema) and not _is_pydantic_model(func_or_schema):
            # Pattern 1: @pipeline (direct decoration, no parentheses)
            direct_func = func_or_schema  # type: ignore[assignment]
        else:
            # Pattern 2: @pipeline(MySchema)
            schema_class = func_or_schema  # type: ignore[assignment]
    # Pattern 3: @pipeline(name="custom") - both None, just use kwargs

    def decorator(func: PipelineFunc[T]) -> PipelineFunc[T]:
        # Validate it's an async function
        if not inspect.iscoroutinefunction(func):
            raise TypeError(
                f"@pipeline can only decorate async functions. "
                f"'{func.__name__}' is not async."
            )

        # Validate signature has ctx parameter
        sig = inspect.signature(func)
        params = list(sig.parameters.values())

        if not params:
            raise TypeError(
                f"@pipeline function '{func.__name__}' must accept a PipelineContext "
                f"as its first parameter."
            )

        # Validate first parameter type annotation if present
        first_param = params[0]
        annotation = first_param.annotation
        if annotation is not inspect.Parameter.empty:
            # Check if annotation is PipelineContext (the Protocol itself)
            if annotation is not PipelineContext:
                raise TypeError(
                    f"@pipeline function '{func.__name__}' first parameter must accept "
                    f"PipelineContext, got '{annotation.__name__ if hasattr(annotation, '__name__') else annotation}'. "
                    f"Use: async def {func.__name__}(ctx: PipelineContext)"
                )

        @functools.wraps(func)
        async def wrapper(ctx: PipelineContext) -> T:
            return await func(ctx)

        # Determine pipeline name
        pipeline_name = name if name else func.__name__

        # Mark as a pipeline for discovery
        wrapper._is_pipeline = True  # type: ignore[attr-defined]
        wrapper._pipeline_name = pipeline_name  # type: ignore[attr-defined]
        wrapper._pipeline_doc = func.__doc__  # type: ignore[attr-defined]

        # Attach schema class if provided
        if schema_class is not None:
            wrapper._input_schema_class = schema_class  # type: ignore[attr-defined]

        # Register in module registry (reject duplicates)
        if pipeline_name in _pipeline_registry:
            raise ValueError(
                f"Pipeline name '{pipeline_name}' is already registered. "
                f"Each @pipeline must have a unique name."
            )
        _pipeline_registry[pipeline_name] = wrapper

        return wrapper  # type: ignore[return-value]

    # If direct decoration (no parentheses), apply decorator immediately
    if direct_func is not None:
        return decorator(direct_func)

    # Otherwise, return decorator for later use
    return decorator


def get_pipeline_registry() -> dict[str, Callable[..., Coroutine[Any, Any, Any]]]:
    """Get the registry of all discovered pipeline functions.

    Used by the server's module loader to find pipeline functions.
    """
    return _pipeline_registry.copy()


def clear_pipeline_registry() -> None:
    """Clear the pipeline registry.

    Useful for testing or when reloading modules.
    """
    _pipeline_registry.clear()
