"""Group key extractors for fair task queue.

This module provides functional extractors for extracting group IDs from TaskInvocation objects.
Extractors are composable callables that can be chained together.
"""

from __future__ import annotations

from haolib.background.fair.abstract import GroupKeyExtractor, TaskInvocation


def extract_from_kwargs(
    key: str,
    fallback_keys: list[str] | None = None,
) -> GroupKeyExtractor[TaskInvocation]:
    """Create an extractor that extracts group ID from TaskInvocation kwargs.

    Args:
        key: Primary key to look for in kwargs
        fallback_keys: Alternative keys to try if primary not found

    Returns:
        Group key extractor function

    Example:
        ```python
        # Extract from kwargs["user_id"] with fallback to "tenant_id"
        extractor = extract_from_kwargs("user_id", fallback_keys=["tenant_id", "group_id"])

        # Use in queue
        queue = FairTaskQueue(..., group_key_extractor=extractor)
        ```

    """

    def extract(invocation: TaskInvocation) -> str:
        # Try primary key
        if key in invocation.kwargs:
            value = invocation.kwargs[key]
            if value is not None:
                return str(value)

        # Try fallback keys
        for fallback_key in fallback_keys or []:
            if fallback_key in invocation.kwargs:
                value = invocation.kwargs[fallback_key]
                if value is not None:
                    return str(value)

        # No group ID found
        keys_tried = [key, *(fallback_keys or [])]
        error_msg = (
            f"Could not extract group ID from task kwargs. Tried keys: {keys_tried}. Task kwargs: {invocation.kwargs}"
        )
        raise ValueError(error_msg)

    return extract


def extract_from_args(index: int = 0) -> GroupKeyExtractor[TaskInvocation]:
    """Create an extractor that extracts group ID from TaskInvocation args.

    Args:
        index: Index in args to extract from (default: 0)

    Returns:
        Group key extractor function

    Example:
        ```python
        # Extract from first positional argument
        extractor = extract_from_args(0)

        # Use in queue
        queue = FairTaskQueue(..., group_key_extractor=extractor)
        ```

    """

    def extract(invocation: TaskInvocation) -> str:
        if index < len(invocation.args):
            value = invocation.args[index]
            if value is not None:
                return str(value)
        error_msg = f"Could not extract group ID from task args at index {index}. Task args: {invocation.args}"
        raise ValueError(error_msg)

    return extract


def extract_with_default(default_group_id: str) -> GroupKeyExtractor[TaskInvocation]:
    """Create an extractor that always returns a default group ID.

    Args:
        default_group_id: Default group ID to return

    Returns:
        Group key extractor function

    Example:
        ```python
        # Always use "default_group"
        extractor = extract_with_default("default_group")

        # Use in queue
        queue = FairTaskQueue(..., group_key_extractor=extractor)
        ```

    """

    def extract(_invocation: TaskInvocation) -> str:
        return default_group_id

    return extract


def compose_extractors(*extractors: GroupKeyExtractor[TaskInvocation]) -> GroupKeyExtractor[TaskInvocation]:
    """Compose multiple extractors, trying each in order until one succeeds.

    Args:
        *extractors: Extractors to try in order

    Returns:
        Composed group key extractor function

    Example:
        ```python
        # Try kwargs first, then args, then default
        extractor = compose_extractors(
            extract_from_kwargs("user_id", fallback_keys=["tenant_id"]),
            extract_from_args(0),
            extract_with_default("default_group"),
        )

        # Use in queue
        queue = FairTaskQueue(..., group_key_extractor=extractor)
        ```

    """
    if not extractors:
        raise ValueError("At least one extractor must be provided")

    def extract(invocation: TaskInvocation) -> str:
        last_error: ValueError | None = None
        for extractor in extractors:
            try:
                return extractor(invocation)
            except ValueError as e:
                last_error = e
                continue

        # All extractors failed
        if last_error:
            error_msg = f"All extractors failed. Last error: {last_error}"
            raise ValueError(error_msg) from last_error
        raise ValueError("All extractors failed")

    return extract


class TaskInvocationGroupKeyExtractor:
    """Group key extractor class for backward compatibility.

    This class wraps the functional extractors to maintain backward compatibility
    with code that uses the class-based API. The class is callable and can be
    used directly as a GroupKeyExtractor.

    Example:
        ```python
        # Old style (still works)
        extractor = TaskInvocationGroupKeyExtractor(
            group_id_key="user_id",
            fallback_keys=["tenant_id"],
        )

        # New style (recommended)
        extractor = extract_from_kwargs("user_id", fallback_keys=["tenant_id"])
        ```

    """

    def __init__(
        self,
        group_id_key: str = "group_id",
        fallback_keys: list[str] | None = None,
        extract_from_args: bool = False,
        args_index: int | None = None,
    ) -> None:
        """Initialize extractor.

        Args:
            group_id_key: Primary key to look for in kwargs
            fallback_keys: Alternative keys to try if primary not found
            extract_from_args: If True, also check args
            args_index: Index in args to check (if extract_from_args=True)

        """
        self._group_id_key = group_id_key
        self._fallback_keys = fallback_keys or []
        self._extract_from_args = extract_from_args
        self._args_index = args_index

        # Build functional extractor chain
        extractors: list[GroupKeyExtractor[TaskInvocation]] = [
            extract_from_kwargs(group_id_key, self._fallback_keys),
        ]

        if self._extract_from_args:
            # Import here to avoid name shadowing with parameter
            from haolib.background.fair.extractors import extract_from_args as _extract_from_args_func  # noqa: PLC0415

            if self._args_index is not None:
                extractors.append(_extract_from_args_func(self._args_index))
            else:
                extractors.append(_extract_from_args_func(0))

        # Compose extractors if multiple, otherwise use single
        if len(extractors) > 1:
            self._extractor = compose_extractors(*extractors)
        else:
            self._extractor = extractors[0]

    def __call__(self, invocation: TaskInvocation) -> str:
        """Extract group ID from TaskInvocation (makes class callable).

        Args:
            invocation: TaskInvocation object

        Returns:
            Group identifier

        Raises:
            ValueError: If group ID cannot be extracted

        """
        return self._extractor(invocation)

    def extract(self, invocation: TaskInvocation) -> str:
        """Extract group ID from TaskInvocation (backward compatibility method).

        Args:
            invocation: TaskInvocation object

        Returns:
            Group identifier

        Raises:
            ValueError: If group ID cannot be extracted

        """
        return self._extractor(invocation)
