# 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 datetime import datetime
from pathlib import Path
from typing import Any, Literal, NewType, Optional

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: Optional[int]
    next: Optional[pydantic.AnyUrl]
    previous: Optional[pydantic.AnyUrl]
    results: list[dict[str, Any]]


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

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


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

    id: int
    created_at: datetime
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    duration: Optional[int] = None
    status: str
    result: str
    worker: Optional[int] = None
    task_name: str
    task_data: dict[str, Any]
    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


# Backport of pydantic.NonNegativeInt from 1.8.
class NonNegativeInt(pydantic.ConstrainedInt):
    """A non-negative integer."""

    ge = 0


# With pydantic >= 2.1.0, use Annotated[str,
# pydantic.StringConstraints(max_length=255)] instead.  Unfortunately there
# doesn't seem to be a less ugly way to write this with earlier versions of
# pydantic that mypy is happy with.
class StrMaxLength255(pydantic.ConstrainedStr):
    """A string with a maximum length of 255 characters."""

    max_length = 255


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

    size: NonNegativeInt
    checksums: dict[str, StrMaxLength255]
    type: Literal["file"]

    @staticmethod
    def create_from(path: Path) -> "FileRequest":
        """Return a FileRequest for the file path."""
        return FileRequest(
            size=pydantic.parse_obj_as(NonNegativeInt, 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: NonNegativeInt
    checksums: dict[str, StrMaxLength255]
    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: Optional[str] = None
    files: FilesRequestType = FilesRequestType({})
    data: dict[str, Any] = {}
    work_request: Optional[int] = None
    expire_at: Optional[datetime] = 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: Optional[datetime] = 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


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())
