from __future__ import annotations

from datetime import datetime, timezone
from typing import Any


class NexusError(Exception):
    """Base structured exception for all NEXUS runtime failures."""

    def __init__(
        self,
        message: str,
        *,
        agent_id: str | None = None,
        node_id: str | None = None,
        context: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(message)
        self.message = message
        self.agent_id = agent_id
        self.node_id = node_id
        self.timestamp = datetime.now(timezone.utc)
        self.context: dict[str, Any] = context.copy() if context else {}

    def to_dict(self) -> dict[str, Any]:
        """Return a serializable representation of the exception."""
        return {
            "type": self.__class__.__name__,
            "message": self.message,
            "agent_id": self.agent_id,
            "node_id": self.node_id,
            "timestamp": self.timestamp.isoformat(),
            "context": self.context,
        }

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.message})"


class ToolFailure(NexusError):
    """Raised when a tool invocation fails."""

    def __init__(self, message: str, tool_name: str, args: dict[str, Any], **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"tool_name": tool_name, "args": args})
        super().__init__(message, context=context, **kwargs)


class AgentFailure(NexusError):
    """Raised when agent task execution or provider call fails."""

    def __init__(self, message: str, attempt: int, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"attempt": attempt})
        super().__init__(message, context=context, **kwargs)


class BudgetExhausted(NexusError):
    """Raised when token usage exceeds configured hard cutoff."""

    def __init__(self, message: str, used_tokens: int, allocated_tokens: int, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"used_tokens": used_tokens, "allocated_tokens": allocated_tokens})
        super().__init__(message, context=context, **kwargs)


class LoopLimitReached(NexusError):
    """Raised when a node loop break condition is met."""

    def __init__(self, message: str, node_id: str, iteration_count: int, strategy: str, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update(
            {
                "node_id": node_id,
                "iteration_count": iteration_count,
                "strategy": strategy,
            }
        )
        super().__init__(message, node_id=node_id, context=context, **kwargs)


class ContextOverflow(NexusError):
    """Raised when a model context window would be exceeded."""

    def __init__(self, message: str, size: int, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"size": size})
        super().__init__(message, context=context, **kwargs)


class OrchestratorError(NexusError):
    """Raised for orchestrator-level failures that are not narrower typed errors."""


class CheckpointError(NexusError):
    """Raised for checkpoint read/write/restore failures."""


class ProviderError(NexusError):
    """Raised when an LLM provider fails or returns invalid output."""

    def __init__(self, message: str, provider_name: str, status_code: int | None = None, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"provider_name": provider_name, "status_code": status_code})
        super().__init__(message, context=context, **kwargs)


class GraphMutationError(NexusError):
    """Raised when an invalid graph mutation is attempted."""


class MCPError(NexusError):
    """Raised for MCP transport and protocol errors."""

    def __init__(self, message: str, server_url: str, **kwargs: Any) -> None:
        context = kwargs.pop("context", {})
        context.update({"server_url": server_url})
        super().__init__(message, context=context, **kwargs)
