# Copyright 2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Models used by debusine tasks."""

from datetime import datetime
from enum import Enum, StrEnum, auto
from itertools import groupby
from typing import Annotated, Any, Literal, TypeAlias, Union

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

import yaml


# taken from functional.py in python3-django 3:4.2.11-1 to avoid the dependency
class classproperty:
    """
    classproperty definition.

    Decorator that converts a method with a single cls argument into a property
    that can be accessed directly from the class.
    """

    def __init__(self, method=None):
        """Define __init__."""
        self.fget = method

    def __get__(self, instance, cls=None):  # noqa:
        """Define __get__."""
        return self.fget(cls)


class TaskTypes(StrEnum):
    """Possible values for task_type."""

    WORKER = "Worker"
    SERVER = "Server"
    INTERNAL = "Internal"
    WORKFLOW = "Workflow"

    @classproperty
    def choices(cls):
        """To be compatible with django.db.models.TextChoice."""
        return [(item, item) for item in cls]


class WorkerType(StrEnum):
    """The type of a Worker: external or Celery."""

    EXTERNAL = "external"
    CELERY = "celery"


class BackendType(StrEnum):
    """Possible values for backend."""

    AUTO = "auto"
    UNSHARE = "unshare"
    SCHROOT = "schroot"
    INCUS_LXC = "incus-lxc"
    INCUS_VM = "incus-vm"
    QEMU = "qemu"


class AutopkgtestNeedsInternet(StrEnum):
    """Possible values for needs_internet."""

    RUN = "run"
    TRY = "try"
    SKIP = "skip"


class LintianPackageType(Enum):
    """Possible package types."""

    SOURCE = auto()
    BINARY_ALL = auto()
    BINARY_ANY = auto()


class LintianFailOnSeverity(StrEnum):
    """Possible values for fail_on_severity."""

    ERROR = "error"
    WARNING = "warning"
    INFO = "info"
    PEDANTIC = "pedantic"
    EXPERIMENTAL = "experimental"
    OVERRIDDEN = "overridden"
    NONE = "none"


class BaseTaskDataModel(pydantic.BaseModel):
    """Stricter pydantic defaults for task data models."""

    class Config:
        """Set up stricter pydantic Config."""

        validate_assignment = True
        extra = pydantic.Extra.forbid

    def dict(self, **kwargs: Any) -> dict[str, Any]:
        """Use aliases by default when serializing."""
        kwargs.setdefault("by_alias", True)
        return super().dict(**kwargs)


# Lookups resulting in exactly one collection item.
LookupSingle: TypeAlias = int | str


def build_lookup_string_segments(*segments: str) -> str:
    """Build a string lookup from segments."""
    return "/".join(segments)


def parse_lookup_string_segments(lookup: str) -> list[str]:
    """Parse a string lookup into segments."""
    return lookup.split("/")


def build_key_value_lookup_segment(
    lookup_type: str, filters: dict[str, str]
) -> str:
    """Build a lookup segment consisting of a series of `key=value` filters."""
    filters_string = ":".join(
        f"{key}={value}" for key, value in filters.items()
    )
    return f"{lookup_type}:{filters_string}"


def parse_key_value_lookup_segment(segment: str) -> tuple[str, dict[str, str]]:
    """Parse a lookup segment consisting of a series of `key=value` filters."""
    # TODO: This is currently only used by the debian:environments
    # collection.  If we generalize this, then we should consider whether we
    # need to add some kind of quoting capability.
    lookup_type, *filters_segments = segment.split(":")
    return lookup_type, dict(item.split("=", 1) for item in filters_segments)


class LookupChildType(StrEnum):
    """Possible values for LookupDict.child_type."""

    BARE = "bare"
    ARTIFACT = "artifact"
    COLLECTION = "collection"


class CollectionItemMatcherKind(StrEnum):
    """Possible values for CollectionItemMatcher.kind."""

    CONTAINS = "contains"
    ENDSWITH = "endswith"
    EXACT = "exact"
    STARTSWITH = "startswith"


class CollectionItemMatcher(BaseTaskDataModel):
    """A matcher for collection item name or per-item data fields."""

    kind: CollectionItemMatcherKind
    value: str


class LookupDict(BaseTaskDataModel):
    """Dictionary lookups for collection items."""

    class Config(BaseTaskDataModel.Config):
        frozen = True

    collection: LookupSingle
    child_type: LookupChildType = LookupChildType.ARTIFACT
    category: str | None = None
    name_matcher: CollectionItemMatcher | None = None
    # Logically a mapping, but we turn it into a tuple so that lookups are
    # hashable.
    data_matchers: tuple[tuple[str, CollectionItemMatcher], ...] = ()

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.root_validator(pre=True, allow_reuse=True)
    def normalize_matchers(
        cls, values: dict[str, Any]  # noqa: U100
    ) -> dict[str, Any]:
        """
        Transform the lookup syntax into a form more convenient for pydantic.

        `name`, `name__*`, `data__KEY`, and `data__KEY__*` keys in `values`
        are transformed into :py:class:`CollectionItemMatcher` instances
        stored in `name_matcher` and `data_matchers`.  Conflicting lookup
        suffixes are rejected.
        """

        def split_words(s: str) -> list[str]:
            return s.split("__")

        def matcher_prefix(s: str) -> tuple[str, ...]:
            match split_words(s):
                case ["name", *_]:
                    return ("name",)
                case ["data", data_key, *_]:
                    return ("data", data_key)
                case _:
                    return ()

        data_matchers: dict[str, CollectionItemMatcher] = dict(
            values.get("data_matchers", ())
        )

        for key, group in groupby(
            sorted(values, key=matcher_prefix), matcher_prefix
        ):
            if not key:
                # Only keys with `name` or `data` as their first segment
                # need transformation.
                continue
            matcher_names = list(group)
            if len(matcher_names) > 1:
                raise ValueError(
                    f"Conflicting matchers: {sorted(matcher_names)}"
                )
            assert len(matcher_names) == 1
            matcher_name = matcher_names[0]

            data_key: str | None = None
            match (words := split_words(matcher_name)):
                case ["name"]:
                    kind = "exact"
                case ["name", kind] if kind != "exact":
                    pass
                case ["data", data_key]:
                    kind = "exact"
                case ["data", data_key, kind] if kind != "exact":
                    pass
                case _:
                    # We filter out keys with first segments other than
                    # `name` and `data` earlier, so anything that reaches
                    # here must be a malformed lookup.
                    raise ValueError(f"Unrecognized matcher: {matcher_name}")

            matcher = CollectionItemMatcher(
                kind=CollectionItemMatcherKind(kind), value=values[matcher_name]
            )
            if words[0] == "name":
                values["name_matcher"] = matcher
            else:
                assert words[0] == "data"
                assert data_key is not None
                data_matchers[data_key] = matcher
            for k in matcher_names:
                del values[k]

        if data_matchers:
            values["data_matchers"] = tuple(sorted(data_matchers.items()))

        return values

    def export(self) -> dict[str, Any]:
        """
        Export the usual input representation of this lookup.

        This reverses the transformations applied by
        :py:meth:`normalize_matchers`.
        """
        value = {"collection": self.collection}
        if self.child_type != LookupChildType.ARTIFACT:
            value["child_type"] = str(self.child_type)
        if self.category is not None:
            value["category"] = self.category
        if self.name_matcher is not None:
            if self.name_matcher.kind == CollectionItemMatcherKind.EXACT:
                value["name"] = self.name_matcher.value
            else:
                value[f"name__{self.name_matcher.kind}"] = (
                    self.name_matcher.value
                )
        for key, matcher in self.data_matchers:
            if matcher.kind == CollectionItemMatcherKind.EXACT:
                value[f"data__{key}"] = matcher.value
            else:
                value[f"data__{key}__{matcher.kind}"] = matcher.value
        return value


class LookupMultiple(BaseTaskDataModel):
    """Lookups resulting in multiple collection items."""

    class Config(BaseTaskDataModel.Config):
        frozen = True

    # LookupMultiple parses a list instead of an object/dict; we turn this
    # into a tuple so that lookups are hashable
    __root__: tuple[LookupSingle | LookupDict, ...]

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.validator("__root__", pre=True, allow_reuse=True)
    def normalize(cls, values: Any) -> tuple[Any, ...]:  # noqa: U100
        """Normalize into a list of multiple matchers."""
        if isinstance(values, dict):
            return (values,)
        elif isinstance(values, (tuple, list)):
            return tuple(values)
        else:
            raise ValueError(
                "Lookup of multiple collection items must be a dictionary or "
                "a list"
            )

    def __iter__(self):
        """Iterate over individual lookups."""
        return iter(self.__root__)

    def export(self) -> dict[str, Any] | list[int | str | dict[str, Any]]:
        """
        Export the usual input representation of this lookup.

        This reverses the transformations applied by :py:meth:`normalize`.
        """
        if len(self.__root__) == 1 and isinstance(self.__root__[0], LookupDict):
            return self.__root__[0].export()
        else:
            return [
                lookup.export() if isinstance(lookup, LookupDict) else lookup
                for lookup in self.__root__
            ]


class NotificationDataEmail(BaseTaskDataModel):
    """Channel data for email notifications."""

    from_: pydantic.EmailStr | None = pydantic.Field(default=None, alias="from")
    to: list[pydantic.EmailStr] | None
    cc: list[pydantic.EmailStr] = pydantic.Field(default_factory=list)
    subject: str | None = None


class ActionTypes(StrEnum):
    """Possible values for EventReaction actions."""

    SEND_NOTIFICATION = "send-notification"
    UPDATE_COLLECTION_WITH_ARTIFACTS = "update-collection-with-artifacts"


class ActionSendNotification(BaseTaskDataModel):
    """structure for notifications."""

    action: Literal[ActionTypes.SEND_NOTIFICATION] = (
        ActionTypes.SEND_NOTIFICATION
    )
    channel: str
    data: NotificationDataEmail | None = None


class ActionUpdateCollectionWithArtifacts(BaseTaskDataModel):
    """structure for collection update."""

    action: Literal[ActionTypes.UPDATE_COLLECTION_WITH_ARTIFACTS] = (
        ActionTypes.UPDATE_COLLECTION_WITH_ARTIFACTS
    )
    collection: LookupSingle
    name_template: str | None = None
    variables: dict[str, str] | None = None
    artifact_filters: dict[str, Any]


EventReaction = Annotated[
    Union[ActionSendNotification, ActionUpdateCollectionWithArtifacts],
    pydantic.Field(discriminator="action"),
]


class EventReactions(BaseTaskDataModel):
    """Structure for event reactions."""

    on_success: list[EventReaction] = []
    on_failure: list[EventReaction] = []


class BaseTaskData(BaseTaskDataModel):
    """
    Base class for task data.

    Task data is encoded as JSON in the database and in the API, and it is
    modeled as a pydantic data structure in memory for both ease of access and
    validation.
    """

    pass


class BaseTaskDataWithExecutor(BaseTaskData):
    """Base task data with fields used to configure executors."""

    backend: BackendType = BackendType.AUTO
    environment: LookupSingle | None = None


class BaseDynamicTaskData(BaseTaskDataModel):
    """
    Base class for dynamic task data.

    This is computed by the scheduler when dispatching a task to a worker.
    It may involve resolving artifact lookups.
    """


class BaseDynamicTaskDataWithExecutor(BaseDynamicTaskData):
    """Dynamic task data for executors."""

    environment_id: int | None = None


class NoopData(BaseTaskData):
    """In memory task data for the Noop task."""

    result: bool = True


def empty_lookup_multiple() -> LookupMultiple:
    """Return an empty :py:class:`LookupMultiple`."""
    return LookupMultiple.parse_obj(())


class AutopkgtestInput(BaseTaskDataModel):
    """Input for an autopkgtest task."""

    source_artifact: LookupSingle
    binary_artifacts: LookupMultiple
    context_artifacts: LookupMultiple = pydantic.Field(
        default_factory=empty_lookup_multiple
    )


class AutopkgtestFailOn(BaseTaskDataModel):
    """Possible values for fail_on."""

    failed_test: bool = True
    flaky_test: bool = False
    skipped_test: bool = False


class AutopkgtestTimeout(BaseTaskDataModel):
    """Timeout specifications for an autopkgtest task."""

    global_: int | None = pydantic.Field(alias="global", ge=0)
    factor: int | None = pydantic.Field(ge=0)
    short: int | None = pydantic.Field(ge=0)
    install: int | None = pydantic.Field(ge=0)
    test: int | None = pydantic.Field(ge=0)
    copy_: int | None = pydantic.Field(alias="copy", ge=0)


class AutopkgtestData(BaseTaskDataWithExecutor):
    """In memory task data for the Autopkgtest task."""

    input: AutopkgtestInput
    host_architecture: str
    include_tests: list[str] = pydantic.Field(default_factory=list)
    exclude_tests: list[str] = pydantic.Field(default_factory=list)
    debug_level: int = pydantic.Field(default=0, ge=0, le=3)
    extra_apt_sources: list[str] = pydantic.Field(default_factory=list)
    use_packages_from_base_repository: bool = False
    extra_environment: dict[str, str] = pydantic.Field(default_factory=dict)
    needs_internet: AutopkgtestNeedsInternet = AutopkgtestNeedsInternet.RUN
    fail_on: AutopkgtestFailOn = pydantic.Field(
        default_factory=AutopkgtestFailOn
    )
    timeout: AutopkgtestTimeout | None
    # BaseTaskDataWithExecutor declares this as optional, but it's required
    # here.
    environment: LookupSingle


class AutopkgtestDynamicData(BaseDynamicTaskDataWithExecutor):
    """Dynamic data for the Autopkgtest task."""

    # BaseDynamicTaskDataWithExecutor declares this as optional, but it's
    # required here.
    environment_id: int

    input_source_artifact_id: int
    input_binary_artifacts_ids: list[int]
    input_context_artifacts_ids: list[int] = pydantic.Field(
        default_factory=list
    )


class LintianInput(BaseTaskDataModel):
    """Input for a lintian task."""

    source_artifact: LookupSingle | None = None
    binary_artifacts: LookupMultiple = pydantic.Field(
        default_factory=empty_lookup_multiple
    )

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.root_validator(allow_reuse=True)
    def check_one_of_source_or_binary(cls, values):  # noqa: U100
        """Ensure a source or binary artifact is present."""
        if values.get("source_artifact") is None and (
            not isinstance(values.get("binary_artifacts"), LookupMultiple)
            or not values.get("binary_artifacts").__root__
        ):
            raise ValueError(
                'One of source_artifact or binary_artifacts must be set'
            )
        return values


class LintianOutput(BaseTaskDataModel):
    """Output configuration for a Lintian task."""

    source_analysis: bool = True
    binary_all_analysis: bool = True
    binary_any_analysis: bool = True


class LintianData(BaseTaskDataWithExecutor):
    """In memory task data for the Lintian task."""

    input: LintianInput
    output: LintianOutput = pydantic.Field(default_factory=LintianOutput)
    target_distribution: str = "debian:unstable"
    # Passed to --tags to lintian
    include_tags: list[str] = pydantic.Field(default_factory=list)
    # Passed to --suppress-tags to lintian
    exclude_tags: list[str] = pydantic.Field(default_factory=list)
    # If the analysis emits tags of this severity or higher, the task will
    # return 'failure' instead of 'success'
    fail_on_severity: LintianFailOnSeverity = LintianFailOnSeverity.NONE


class LintianDynamicData(BaseDynamicTaskDataWithExecutor):
    """Dynamic data for the Lintian task."""

    input_source_artifact_id: int | None = None
    input_binary_artifacts_ids: list[int] = pydantic.Field(default_factory=list)


class BlhcInput(BaseTaskDataModel):
    """Input for a blhc task."""

    artifact: LookupSingle


class BlhcOutput(BaseTaskDataModel):
    """Output configuration for a blhc task."""

    source_analysis: bool = True


class BlhcFlags(StrEnum):
    """Possible values for extra_flags."""

    ALL = "--all"
    BINDNOW = "--bindnow"
    BUILDD = "--buildd"
    COLOR = "--color"
    DEBIAN = "--debian"
    LINE_NUMBERS = "--line-numbers"
    PIE = "--pie"


class BlhcData(BaseTaskDataWithExecutor):
    """In memory task data for the Blhc task."""

    input: BlhcInput
    output: BlhcOutput = pydantic.Field(default_factory=BlhcOutput)
    # Passed to blhc
    extra_flags: list[BlhcFlags] = pydantic.Field(default_factory=list)


class BlhcDynamicData(BaseDynamicTaskDataWithExecutor):
    """Dynamic data for the Blhc task."""

    input_artifact_id: int


class MmDebstrapVariant(StrEnum):
    """Possible values for variant."""

    BUILDD = "buildd"
    MINBASE = "minbase"
    DASH = "-"
    APT = "apt"
    CUSTOM = "custom"
    DEBOOTSTRAP = "debootstrap"
    ESSENTIAL = "essential"
    EXTRACT = "extract"
    IMPORTANT = "important"
    REQUIRED = "required"
    STANDARD = "standard"


class SystemBootstrapOptionsVariant(StrEnum):
    """Possible values for system bootstrap variant."""

    BUILDD = "buildd"
    MINBASE = "minbase"


# Currently needed by SimpleSystemImageBuild, workaround for
# https://github.com/yaml/pyyaml/issues/722
yaml.SafeDumper.add_representer(
    SystemBootstrapOptionsVariant,
    yaml.representer.SafeRepresenter.represent_str,
)


class SystemBootstrapOptions(BaseTaskDataModel):
    """Structure of SystemBootstrap options."""

    architecture: str
    variant: SystemBootstrapOptionsVariant | None = None
    extra_packages: list[str] = pydantic.Field(default_factory=list)


class SystemBootstrapRepositoryType(StrEnum):
    """Possible values for repository types."""

    DEB = "deb"
    DEB_SRC = "deb-src"


class SystemBootstrapRepositoryCheckSignatureWith(StrEnum):
    """Possible values for check_signature_with."""

    SYSTEM = "system"
    EXTERNAL = "external"
    NO_CHECK = "no-check"


class SystemBootstrapRepositoryKeyring(BaseTaskDataModel):
    """Description of a repository keyring."""

    url: pydantic.AnyUrl
    sha256sum: str = ""
    install: bool = False


class SystemBootstrapRepository(BaseTaskDataModel):
    """Description of one repository in SystemBootstrapData."""

    mirror: str
    suite: str
    types: list[SystemBootstrapRepositoryType] = pydantic.Field(
        default_factory=lambda: [SystemBootstrapRepositoryType.DEB],
        min_items=1,
        unique_items=True,
    )
    components: list[str] | None = pydantic.Field(None, unique_items=True)
    check_signature_with: SystemBootstrapRepositoryCheckSignatureWith = (
        SystemBootstrapRepositoryCheckSignatureWith.SYSTEM
    )
    keyring_package: str | None = None
    keyring: SystemBootstrapRepositoryKeyring | None = None

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.root_validator(allow_reuse=True)
    def _check_external_keyring(cls, values: Any) -> Any:  # noqa: U100
        """Require keyring if check_signature_with is external."""
        if (
            values.get("check_signature_with")
            == SystemBootstrapRepositoryCheckSignatureWith.EXTERNAL
        ):
            if values.get("keyring") is None:
                raise ValueError(
                    "repository requires 'keyring': "
                    "'check_signature_with' is set to 'external'"
                )
        return values


class SystemBootstrapData(BaseTaskData):
    """Base for in-memory class data for SystemBootstrap tasks."""

    bootstrap_options: SystemBootstrapOptions
    bootstrap_repositories: list[SystemBootstrapRepository] = pydantic.Field(
        min_items=1
    )
    customization_script: str | None = None


class MmDebstrapBootstrapOptions(BaseTaskDataModel):
    """Structure of MmDebstrap options."""

    architecture: str
    # TODO: enums cannot inherit from other enums: find out how to make this
    # extension type safe
    variant: MmDebstrapVariant | None = None
    extra_packages: list[str] = pydantic.Field(default_factory=list)
    use_signed_by: bool = True


class MmDebstrapData(BaseTaskData):
    """In memory task data for the MmDebstrap task."""

    bootstrap_options: MmDebstrapBootstrapOptions
    bootstrap_repositories: list[SystemBootstrapRepository] = pydantic.Field(
        min_items=1
    )
    customization_script: str | None = None


class DiskImageFormat(StrEnum):
    """Possible disk image formats."""

    RAW = "raw"
    QCOW2 = "qcow2"


class Partition(BaseTaskDataModel):
    """Partition definition."""

    size: int
    filesystem: str
    mountpoint: str = "none"


class DiskImage(BaseTaskDataModel):
    """Disk image definition."""

    format: DiskImageFormat
    filename: str = "image"
    kernel_package: str | None = None
    bootloader: str | None = None
    partitions: list[Partition] = pydantic.Field(min_items=1)


class SystemImageBuildData(BaseTaskData):
    """Base for in-memory class data for SystemImageBuild tasks."""

    bootstrap_options: SystemBootstrapOptions
    bootstrap_repositories: list[SystemBootstrapRepository]
    customization_script: str | None = None
    disk_image: DiskImage


class PiupartsDataInput(BaseTaskDataModel):
    """Input for a piuparts task."""

    binary_artifacts: LookupMultiple


class PiupartsData(BaseTaskDataWithExecutor):
    """In memory task data for the Piuparts task."""

    input: PiupartsDataInput
    host_architecture: str
    base_tgz: LookupSingle
    # BaseTaskDataWithExecutor declares this as optional, but it's required
    # here.
    environment: LookupSingle


class PiupartsDynamicData(BaseDynamicTaskDataWithExecutor):
    """Dynamic data for the Piuparts task."""

    # BaseTaskDynamicDataWithExecutor declares this as optional, but it's
    # required here.
    environment_id: int

    input_binary_artifacts_ids: list[int]
    base_tgz_id: int


class SbuildInput(BaseTaskDataModel):
    """Input for a sbuild task."""

    source_artifact: LookupSingle
    extra_binary_artifacts: LookupMultiple = pydantic.Field(
        default_factory=empty_lookup_multiple
    )


class SbuildBuildComponent(StrEnum):
    """Possible values for build_components."""

    ANY = "any"
    ALL = "all"
    SOURCE = "source"


class SbuildData(BaseTaskDataWithExecutor):
    """In memory task data for the Sbuild task."""

    input: SbuildInput
    host_architecture: str
    distribution: str | None = None
    build_components: list[SbuildBuildComponent] = pydantic.Field(
        default_factory=lambda: [SbuildBuildComponent.ANY],
    )

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.root_validator
    def check_one_of_backend_or_environment(cls, values):  # noqa: U100
        """Make sure that backends have the arguments they need."""
        if values.get("backend") == BackendType.SCHROOT:
            if values.get("distribution") is None:
                raise ValueError(
                    'The backend "schroot" requires "distribution" to be set'
                )
            if values.get("environment") is not None:
                raise ValueError(
                    'The backend "schroot" requires "environment" to be unset'
                )
        else:
            if values.get("environment") is None:
                raise ValueError(
                    f'''The backend "{values.get('backend')}"'''
                    ' requires "environment" to be set'
                )
            if values.get("distribution") is not None:
                raise ValueError(
                    f'''The backend "{values.get('backend')}"'''
                    ' requires "distribution" to be unset'
                )
        return values


class SbuildDynamicData(BaseDynamicTaskDataWithExecutor):
    """Dynamic data for the Sbuild task."""

    input_source_artifact_id: int
    input_extra_binary_artifacts_ids: list[int] = pydantic.Field(
        default_factory=list
    )


class ImageCacheUsageLogEntry(BaseTaskDataModel):
    """Entry in ImageCacheUsageLog for cached executor images."""

    filename: str
    backend: str | None = None
    timestamp: datetime

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.validator("timestamp")
    def timestamp_is_aware(cls, timestamp: datetime) -> datetime:  # noqa: U100
        """Ensure that the timestamp is TZ-aware."""
        tzinfo = timestamp.tzinfo
        if tzinfo is None or tzinfo.utcoffset(timestamp) is None:
            raise ValueError("timestamp is TZ-naive")
        return timestamp


class ImageCacheUsageLog(BaseTaskDataModel):
    """Usage log for cached executor images."""

    version: int = 1
    backends: set[str] = pydantic.Field(default_factory=set)
    usage: list[ImageCacheUsageLogEntry] = pydantic.Field(default_factory=list)

    # FIXME: pydantic 1 validators are class methods, and pydantic mandates the
    # first argument to be `cls`; however calling it anything other than `self`
    # makes flake8 complain about an unused argument
    @pydantic.validator("version")
    def version_is_known(cls, version: int) -> int:  # noqa: U100
        """Ensure that the version is known."""
        if version != 1:
            raise ValueError(f"Unknown usage log version {version}")
        return version
