from dataclasses import dataclass
from decimal import Decimal
from typing import Optional, Sequence

from aiohttp import web

from ai.backend.common.exception import (
    BackendAIError,
    ErrorCode,
    ErrorDetail,
    ErrorDomain,
    ErrorOperation,
)
from ai.backend.common.types import DeviceId, SlotName


class InitializationError(Exception):
    """
    Errors during agent initialization and compute plugin setup
    """

    pass


class ResourceError(ValueError):
    pass


class InvalidArgumentError(RuntimeError):
    pass


class UnsupportedResource(ResourceError):
    pass


class InvalidResourceCombination(ResourceError):
    pass


class InvalidResourceArgument(ResourceError):
    pass


class NotMultipleOfQuantum(InvalidResourceArgument):
    pass


@dataclass
class InsufficientResource(ResourceError):
    msg: str
    slot_name: SlotName
    requested_alloc: Decimal
    total_allocatable: Decimal | int
    allocation: dict[SlotName, dict[DeviceId, Decimal]]
    context_tag: Optional[str] = None

    def __str__(self) -> str:
        return (
            f"InsufficientResource: {self.msg} ({self.slot_name}"
            + (f" (tag: {self.context_tag!r}), " if self.context_tag else ", ")
            + f"allocating {self.requested_alloc} out of {self.total_allocatable})"
        )

    def __reduce__(self):
        return (
            self.__class__,
            (
                self.msg,
                self.slot_name,
                self.requested_alloc,
                self.total_allocatable,
                self.allocation,
                self.context_tag,
            ),
        )


@dataclass
class FractionalResourceFragmented(ResourceError):
    msg: str
    slot_name: SlotName
    requested_alloc: Decimal
    dev_allocs: Sequence[tuple[DeviceId, Decimal]]
    context_tag: Optional[str] = None

    def __str__(self) -> str:
        return (
            f"FractionalResourceFragmented: {self.msg} ({self.slot_name}"
            + (f" (tag: {self.context_tag!r}), " if self.context_tag else ", ")
            + f"allocating {self.requested_alloc} from {self.dev_allocs})"
        )

    def __reduce__(self):
        return (
            self.__class__,
            (
                self.msg,
                self.slot_name,
                self.requested_alloc,
                self.dev_allocs,
                self.context_tag,
            ),
        )


class UnsupportedBaseDistroError(RuntimeError):
    pass


class ContainerCreationError(Exception):
    container_id: str

    def __init__(self, container_id: str, message: str | None = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.container_id = container_id
        self.message = message


class K8sError(Exception):
    def __init__(self, message):
        super().__init__(message)
        self.message = message


class AgentError(RuntimeError):
    """
    A dummy exception class to distinguish agent-side errors passed via
    aiozmq.rpc calls.

    It carrise two args tuple: the exception type and exception arguments from
    the agent.
    """

    def __init__(self, *args, exc_repr: Optional[str] = None):
        super().__init__(*args)
        self.exc_repr = exc_repr


class InvalidSocket(Exception):
    pass


class NetworkPluginNotFound(RuntimeError):
    pass


class ServicePortAlreadyUsedError(RuntimeError):
    """
    Raised when a port is already used by another service.
    """


class InvalidModelConfigurationError(RuntimeError):
    """
    Raised when a model configuration is invalid.
    """


class KernelRegistryNotFound(BackendAIError, web.HTTPNotFound):
    error_type = "https://backend.ai/errors/agent/kernel-registry-not-found"
    error_title = "Kernel Registry Not Found"

    @classmethod
    def error_code(cls) -> ErrorCode:
        return ErrorCode(
            domain=ErrorDomain.KERNEL_REGISTRY,
            operation=ErrorOperation.READ,
            error_detail=ErrorDetail.NOT_FOUND,
        )
