"""
Policies — structured error handling and recovery strategies.

Design notes:
  - Three-tier policy system:
    1. RetryPolicy: retry with the same agent (backoff, max attempts)
    2. FallbackPolicy: switch to an alternative agent
    3. EscalationPolicy: escalate to a human or supervisor agent
  - Policies are composable: try retry → fallback → escalation
  - Each policy is a simple callable for maximum flexibility
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger("nexus.runtime.policies")


# Type alias for the task execution function
TaskExecutor = Callable[..., Awaitable[Any]]


@dataclass
class RetryPolicy:
    """
    Retry a failed task with the same agent.

    Supports:
    - Fixed or exponential backoff
    - Max attempts
    - Retry condition (which exceptions to retry on)
    """

    max_attempts: int = 3
    initial_delay: float = 1.0  # seconds
    max_delay: float = 60.0
    backoff_factor: float = 2.0
    retry_on: tuple[type[Exception], ...] = (Exception,)

    async def execute(
        self,
        func: TaskExecutor,
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        """
        Execute a function with retry logic.

        Returns the function result on success, or raises the last exception.
        """
        last_exception: Exception | None = None
        delay = self.initial_delay

        for attempt in range(1, self.max_attempts + 1):
            try:
                result = await func(*args, **kwargs)
                if attempt > 1:
                    logger.info(f"Retry succeeded on attempt {attempt}")
                return result
            except self.retry_on as exc:
                last_exception = exc
                if attempt < self.max_attempts:
                    logger.warning(
                        f"Attempt {attempt}/{self.max_attempts} failed: {exc}. "
                        f"Retrying in {delay:.1f}s..."
                    )
                    await asyncio.sleep(delay)
                    delay = min(delay * self.backoff_factor, self.max_delay)
                else:
                    logger.error(
                        f"All {self.max_attempts} attempts failed. Last error: {exc}"
                    )

        if last_exception:
            raise last_exception
        raise RuntimeError("RetryPolicy exhausted without result or exception")


@dataclass
class FallbackPolicy:
    """
    Fall back to an alternative agent when the primary agent fails.

    The fallback chain is tried in order. If all fallbacks fail,
    the original exception is raised.
    """

    fallback_agent_ids: list[str] = field(default_factory=list)
    # Optional: function to select the best fallback dynamically
    fallback_selector: Callable[[str, Exception], str | None] | None = None

    async def execute(
        self,
        primary_func: TaskExecutor,
        fallback_func: Callable[[str], TaskExecutor],
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        """
        Try the primary function, then fall back to alternatives.

        Args:
            primary_func: The main execution function.
            fallback_func: A factory that creates an executor for a given agent_id.
            *args, **kwargs: Passed to the executor.
        """
        try:
            return await primary_func(*args, **kwargs)
        except Exception as primary_exc:
            logger.warning(f"Primary agent failed: {primary_exc}, trying fallbacks")

            # Try each fallback agent
            for agent_id in self.fallback_agent_ids:
                try:
                    executor = fallback_func(agent_id)
                    result = await executor(*args, **kwargs)
                    logger.info(f"Fallback to agent {agent_id} succeeded")
                    return result
                except Exception as fallback_exc:
                    logger.warning(
                        f"Fallback agent {agent_id} also failed: {fallback_exc}"
                    )

            # Dynamic fallback selection
            if self.fallback_selector:
                selected = self.fallback_selector("", primary_exc)
                if selected:
                    try:
                        executor = fallback_func(selected)
                        return await executor(*args, **kwargs)
                    except Exception:
                        pass

            # All fallbacks exhausted
            raise primary_exc


@dataclass
class EscalationPolicy:
    """
    Escalate to a human or supervisor agent when automated handling fails.

    This is the last resort in the policy chain.
    """

    escalation_handler: Callable[[str, str, Exception], Awaitable[Any]] | None = None
    auto_escalate_on: tuple[type[Exception], ...] = ()
    max_wait_seconds: float = 300.0  # Max time to wait for human response

    async def escalate(
        self,
        task_id: str,
        agent_id: str,
        error: Exception,
    ) -> Any:
        """
        Escalate a failed task to a human or supervisor.

        Returns the human/supervisor's response, or raises if no handler is configured.
        """
        logger.warning(
            f"Escalating task {task_id} (agent {agent_id}): {error}"
        )

        if self.escalation_handler:
            return await self.escalation_handler(task_id, agent_id, error)

        raise RuntimeError(
            f"Task {task_id} requires human intervention but no escalation_handler "
            f"is configured. Error: {error}"
        )


@dataclass
class PolicyChain:
    """
    Compose retry, fallback, and escalation policies into a chain.

    The chain tries each policy in order:
    1. RetryPolicy (same agent, multiple attempts)
    2. FallbackPolicy (different agent)
    3. EscalationPolicy (human in the loop)
    """

    retry: RetryPolicy = field(default_factory=RetryPolicy)
    fallback: FallbackPolicy = field(default_factory=FallbackPolicy)
    escalation: EscalationPolicy = field(default_factory=EscalationPolicy)

    async def execute(
        self,
        task_id: str,
        agent_id: str,
        primary_func: TaskExecutor,
        fallback_func: Callable[[str], TaskExecutor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        """
        Execute with the full policy chain: retry → fallback → escalate.
        """
        last_error: Exception | None = None

        # Step 1: Retry with the same agent
        try:
            return await self.retry.execute(primary_func, *args, **kwargs)
        except Exception as retry_exc:
            last_error = retry_exc
            logger.warning(f"Retry policy exhausted for task {task_id}: {retry_exc}")

        # Step 2: Fallback to alternative agents
        if fallback_func and self.fallback.fallback_agent_ids:
            try:
                return await self.fallback.execute(
                    primary_func, fallback_func, *args, **kwargs
                )
            except Exception as fallback_exc:
                last_error = fallback_exc
                logger.warning(
                    f"Fallback policy exhausted for task {task_id}: {fallback_exc}"
                )

        # Step 3: Escalate to human
        assert last_error is not None
        try:
            return await self.escalation.escalate(task_id, agent_id, last_error)
        except Exception:
            raise last_error  # Re-raise the original error if escalation also fails
