# Copyright 2022 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 client."""

import json
from collections.abc import Sequence
from datetime import datetime
from enum import StrEnum
from pathlib import Path
from typing import Annotated, Any, Generic, Literal, NewType, TypeVar

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

from debusine.utils import calculate_hash


class StrictBaseModel(pydantic.BaseModel):
    """Stricter pydantic configuration."""

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

        validate_assignment = True


class PaginatedResponse(StrictBaseModel):
    """Paginated response from the API."""

    count: int | None
    next: pydantic.AnyUrl | None
    previous: pydantic.AnyUrl | None
    results: list[dict[str, Any]]


class WorkRequestRequest(StrictBaseModel):
    """Client send a WorkRequest to the server."""

    task_name: str
    workspace: str | None = None
    task_data: dict[str, Any]
    event_reactions: dict[str, Any]


class WorkRequestResponse(StrictBaseModel):
    """Server return a WorkRequest to the client."""

    id: int
    created_at: datetime
    started_at: datetime | None = None
    completed_at: datetime | None = None
    duration: int | None = None
    status: str
    result: str
    worker: int | None = None
    task_name: str
    task_data: dict[str, Any]
    dynamic_task_data: dict[str, Any] | None
    priority_base: int
    priority_adjustment: int
    artifacts: list[int]
    workspace: str

    def __str__(self):
        """Return representation of the object."""
        return f'WorkRequest: {self.id}'


class OnWorkRequestCompleted(StrictBaseModel):
    """
    Server return an OnWorkRequestCompleted to the client.

    Returned via websocket consumer endpoint.
    """

    work_request_id: int
    completed_at: datetime
    result: str


class WorkflowTemplateRequest(StrictBaseModel):
    """Client sends a WorkflowTemplate to the server."""

    name: str
    task_name: str
    workspace: str | None = None
    task_data: dict[str, Any]
    priority: int


class WorkflowTemplateResponse(StrictBaseModel):
    """Server returns a WorkflowTemplate to the server."""

    id: int
    task_name: str
    workspace: str
    task_data: dict[str, Any]
    priority: int


class CreateWorkflowRequest(StrictBaseModel):
    """Client sends a workflow creation request to the server."""

    template_name: str
    workspace: str | None = None
    task_data: dict[str, Any]


class FileRequest(StrictBaseModel):
    """Declare a FileRequest: client sends it to the server."""

    size: int = pydantic.Field(ge=0)
    checksums: dict[str, Annotated[str, pydantic.Field(max_length=255)]]
    type: Literal["file"]

    @staticmethod
    def create_from(path: Path) -> "FileRequest":
        """Return a FileRequest for the file path."""
        return FileRequest(
            size=path.stat().st_size,
            checksums={"sha256": calculate_hash(path, "sha256").hex()},
            type="file",
        )


class FileResponse(StrictBaseModel):
    """Declare a FileResponse: server sends it to the client."""

    size: int = pydantic.Field(ge=0)
    checksums: dict[str, Annotated[str, pydantic.Field(max_length=255)]]
    type: Literal["file"]
    url: pydantic.AnyUrl


FilesRequestType = NewType("FilesRequestType", dict[str, FileRequest])
FilesResponseType = NewType("FilesResponseType", dict[str, FileResponse])


class ArtifactCreateRequest(StrictBaseModel):
    """Declare an ArtifactCreateRequest: client sends it to the server."""

    category: str
    workspace: str | None = None
    files: FilesRequestType = FilesRequestType({})
    data: dict[str, Any] = {}
    work_request: int | None = None
    expire_at: datetime | None = None


class ArtifactResponse(StrictBaseModel):
    """Declare an ArtifactResponse: server sends it to the client."""

    id: int
    workspace: str
    category: str
    created_at: datetime
    data: dict[str, Any]
    download_tar_gz_url: pydantic.AnyUrl
    files_to_upload: list[str]
    expire_at: datetime | None = None
    files: FilesResponseType = FilesResponseType({})


class RemoteArtifact(StrictBaseModel):
    """Declare RemoteArtifact."""

    id: int
    workspace: str


RelationCreateRequestType = Literal["extends", "relates-to", "built-using"]


class RelationCreateRequest(StrictBaseModel):
    """Declare a RelationCreateRequest: client sends it to the server."""

    artifact: int
    target: int
    type: RelationCreateRequestType


class RelationResponse(RelationCreateRequest):
    """Declare a RelationResponse."""

    id: int


class CollectionItemType(StrEnum):
    """A collection item type."""

    BARE = "b"
    ARTIFACT = "a"
    COLLECTION = "c"


class LookupSingleRequest(StrictBaseModel):
    """A request from the client to look up a single collection item."""

    lookup: int | str
    work_request: int
    expect_type: CollectionItemType
    default_category: str | None = None


class LookupMultipleRequest(StrictBaseModel):
    """A request from the client to look up multiple collection items."""

    lookup: list[int | str | dict[str, Any]]
    work_request: int
    expect_type: CollectionItemType
    default_category: str | None = None


class LookupSingleResponse(StrictBaseModel):
    """A response from the server with a single lookup result."""

    result_type: CollectionItemType
    collection_item: int | None = None
    artifact: int | None = None
    collection: int | None = None


class LookupSingleResponseArtifact(LookupSingleResponse):
    """
    A response from the server with a single lookup result for an artifact.

    Used to assist type annotations.
    """

    result_type: Literal[CollectionItemType.ARTIFACT]
    artifact: int


class LookupSingleResponseCollection(LookupSingleResponse):
    """
    A response from the server with a single lookup result for a collection.

    Used to assist type annotations.
    """

    result_type: Literal[CollectionItemType.COLLECTION]
    collection: int


LSR = TypeVar("LSR", bound=LookupSingleResponse, covariant=True)


class LookupMultipleResponse(StrictBaseModel, Generic[LSR]):
    """A response from the server with multiple lookup results."""

    # This model parses a list, not an object/dict.
    __root__: Sequence[LSR]

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


def model_to_json_serializable_dict(
    model: pydantic.BaseModel,
) -> dict[Any, Any]:
    """
    Similar to model.dict() but the returned dictionary is JSON serializable.

    For example, a datetime() is not JSON serializable. Using this method will
    return a dictionary with a string instead of a datetime object.
    """
    return json.loads(model.json())
