# 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, auto
from typing import Any, Optional, Union


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

import yaml


# Convert enums to StrEnum on python 3.11+


class WorkerType(str, Enum):
    """The type of a Worker: external or Celery."""

    EXTERNAL = "external"
    CELERY = "celery"


class BackendType(str, Enum):
    """Possible values for backend."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


class AutopkgtestNeedsInternet(str, Enum):
    """Possible values for needs_internet."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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

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


# Convert to StrEnum on python 3.11+
class LintianFailOnSeverity(str, Enum):
    """Possible values for fail_on_severity."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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)


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

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


class TaskNotification(BaseTaskDataModel):
    """on_failure structure for task notifications."""

    channel: str
    # Use Union instead of Optional to mark that this can hold multiple kinds
    # of notification channel data
    data: Union[NotificationDataEmail, None] = None


class TaskNotifications(BaseTaskDataModel):
    """Data structure for task notifications."""

    on_failure: list[TaskNotification] = pydantic.Field(min_items=1)


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.
    """

    notifications: Optional[TaskNotifications] = None


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

    backend: BackendType = BackendType.AUTO
    environment_id: Optional[int] = None


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

    result: bool = True


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

    source_artifact_id: int
    binary_artifacts_ids: list[int]
    context_artifacts_ids: list[int] = pydantic.Field(default_factory=list)


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_: Optional[int] = pydantic.Field(alias="global", ge=0)
    factor: Optional[int] = pydantic.Field(ge=0)
    short: Optional[int] = pydantic.Field(ge=0)
    install: Optional[int] = pydantic.Field(ge=0)
    test: Optional[int] = pydantic.Field(ge=0)
    copy_: Optional[int] = 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
    environment: dict[str, str] = pydantic.Field(default_factory=dict)
    needs_internet: AutopkgtestNeedsInternet = AutopkgtestNeedsInternet.RUN
    fail_on: AutopkgtestFailOn = pydantic.Field(
        default_factory=AutopkgtestFailOn
    )
    timeout: Optional[AutopkgtestTimeout]
    # BaseTaskDataWithExecutor declares this as optional, but it's required
    # here.
    environment_id: int


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

    source_artifact_id: Optional[int] = None
    binary_artifacts_ids: list[int] = 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.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_id") is None and (
            values.get("binary_artifacts_ids") is None
            or not bool(values.get("binary_artifacts_ids"))
        ):
            raise ValueError(
                'One of source_artifact_id or binary_artifacts_ids 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 BlhcInput(BaseTaskDataModel):
    """Input for a blhc task."""

    artifact_id: int


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

    source_analysis: bool = True


class BlhcFlags(str, Enum):
    """Possible values for extra_flags."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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 MmDebstrapVariant(str, Enum):
    """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"

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


class SystemBootstrapOptionsVariant(str, Enum):
    """Possible values for system bootstrap variant."""

    BUILDD = "buildd"
    MINBASE = "minbase"

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


# 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: Optional[SystemBootstrapOptionsVariant] = None
    extra_packages: list[str] = pydantic.Field(default_factory=list)


class SystemBootstrapRepositoryType(str, Enum):
    """Possible values for repository types."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


class SystemBootstrapRepositoryCheckSignatureWith(str, Enum):
    """Possible values for check_signature_with."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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: Optional[list[str]] = pydantic.Field(None, unique_items=True)
    check_signature_with: SystemBootstrapRepositoryCheckSignatureWith = (
        SystemBootstrapRepositoryCheckSignatureWith.SYSTEM
    )
    keyring_package: Optional[str] = None
    keyring: Optional[SystemBootstrapRepositoryKeyring] = 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: Optional[str] = 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: Optional[MmDebstrapVariant] = 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: Optional[str] = None


class DiskImageFormat(str, Enum):
    """Possible disk image formats."""

    RAW = "raw"
    QCOW2 = "qcow2"

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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: Optional[str] = None
    bootloader: Optional[str] = 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: Optional[str] = None
    disk_image: DiskImage


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

    binary_artifacts_ids: list[int]


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

    input: PiupartsDataInput
    host_architecture: str
    base_tgz_id: int
    # BaseTaskDataWithExecutor declares this as optional, but it's required
    # here.
    environment_id: int


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

    source_artifact_id: int


class SbuildBuildComponent(str, Enum):
    """Possible values for build_components."""

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

    def __str__(self):
        """Use enum value as string representation."""
        return self.value


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

    input: SbuildInput
    host_architecture: str
    distribution: Optional[str] = None
    build_components: list[SbuildBuildComponent] = pydantic.Field(
        default_factory=lambda: [SbuildBuildComponent.ANY],
    )
    sbuild_options: list[str] = 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.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_id") is not None:
                raise ValueError(
                    'The backend "schroot" requires "environment_id" to be '
                    'unset'
                )
        else:
            if values.get("environment_id") is None:
                raise ValueError(
                    f'''The backend "{values.get('backend')}"'''
                    ' requires "environment_id" 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 ImageCacheUsageLogEntry(BaseTaskDataModel):
    """Entry in ImageCacheUsageLog for cached executor images."""

    filename: str
    backend: Optional[str] = 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
