"""Endpoints requests/responses."""
from __future__ import annotations
from dataclasses import dataclass
import datetime as dt
from typing import Any, Dict, List, Optional, Union

from terality_serde import SerdeMixin, SerializableEnum

from . import TransferConfig, TransferConfigLocal


@dataclass
class PendingComputationResponse(SerdeMixin):
    pending_computation_id: str


@dataclass
class ComputationResponse(SerdeMixin):
    result: Any
    inplace: bool


@dataclass
class CreateSessionResponse(SerdeMixin):
    id: str  # pylint: disable=invalid-name
    upload_config: Union[TransferConfig, TransferConfigLocal]
    download_config: Union[TransferConfig, TransferConfigLocal]


@dataclass
class DeleteSessionResponse(SerdeMixin):
    """Response to the delete session endpoint"""


@dataclass
class Upload(SerdeMixin):
    path: str


@dataclass
class UploadRequest(SerdeMixin):
    path: str
    transfer_id: str
    aws_region: Optional[str]

    def terality_hash(self):
        return self.path


@dataclass
class ExportRequest(SerdeMixin):
    # Parameters of the "to_xxx" (to_csv, to_parquet...) function.
    # The server will echo it back in the response, but (as of now) don't use this information otherwise.
    path: str
    storage_options: Optional[Dict[str, Any]]
    # Preferred region for the Terality export bucket.
    aws_region: Optional[str]


@dataclass
class ExportResponse(SerdeMixin):
    path: str
    storage_options: Optional[Dict[str, Any]]
    aws_region: Optional[str]
    transfer_id: str

    is_folder: bool  # True means it is the result of to_csv_folder or to_parquet_folder
    with_leading_zeros: bool = False  # Only relevant if is_folder is True. Output filename numbers will have leading zeros.


@dataclass
class AwsCredentials(SerdeMixin):
    access_key_id: str
    secret_access_key: str
    session_token: str


@dataclass
class DataTransferResponse(SerdeMixin):
    temporary_upload_aws_credentials: AwsCredentials


class StorageService(SerializableEnum):
    AWS_S3 = "AWS_S3"
    AZURE_BLOB_STORAGE = "AZURE_BLOB_STORAGE"


@dataclass
class AwsPresignedUrlSource(SerdeMixin):
    """A presigned AWS URL allowing to run GetObject on an S3 object part."""

    # Possible optimization: for the same object (bucket + key), the same presigned URL could
    # be reused for each range to get (the same presigned URL can be requested with different
    # Range HTTP headers).
    url: str
    object_key: ObjectStorageKey
    # warning, these offsets are inclusive
    range_start_byte: int
    range_end_byte: int


@dataclass(frozen=True, order=True)
class ObjectStorageKey(SerdeMixin):
    """Generic representation of the path to an object (= blob in Azure) in some cloud storage service.

    (AWS Simple Storage Service (S3), Google Cloud Storage (GCS), Azure Blob Storage...)
    """

    bucket: str  # maps to a container in Azure Blob Storage, called "bucket" in most of the other services
    # (Azure Blob Storage also needs an extra "storage account name" information to uniquely identity
    # the blob. It should be transmitted along this object).

    key: str  # also called an "object name" in Google Cloud Storage and a "blob name" in Azure


@dataclass
class ImportFromS3Source(SerdeMixin):
    presigned_urls: List[AwsPresignedUrlSource]


@dataclass
class AzureBlobSource(SerdeMixin):
    """A reference to an blob stored in Azure Blob storage, plus a shared access signature (SAS) that can read it.

    The SAS is a temporary credential generated by the Terality client that can be used as a credential
    to get read-only access to the given blob.
    """

    object_key: ObjectStorageKey
    shared_access_signature: str
    storage_account_name: str
    # Possible optimization: for the same object (container + name), the same SAS could
    # be reused for each range to get.
    range_start_byte: int
    range_end_byte: int


# Even if the user provided an Azure Datalake Gen2 URL, the server always work
# at the Blob Storage level for import/exports. The client does the mapping.
@dataclass
class ImportFromAzureBlobStorageSource(SerdeMixin):
    objects: List[AzureBlobSource]


@dataclass
class ImportFromCloudRequest(SerdeMixin):
    service: StorageService
    source: Union[ImportFromS3Source, ImportFromAzureBlobStorageSource]


@dataclass
class ImportFromCloudResponse(SerdeMixin):
    transfer_id: str


@dataclass
class AwsS3ObjectPartExportRequest(SerdeMixin):
    """A presigned AWS URL allowing to run UploadPart on a S3 object part."""

    # Source
    source_object_key: ObjectStorageKey
    range_start_byte: int
    range_end_byte: int

    # Destination
    presigned_url: str

    # The server doesn't need to know the following pieces of information, but it makes writing the client
    # easier. Instead of matching the server results back using the URL (or an arbitrary ID), the server can
    # send back these IDs and the client can directly use them to complete the upload.
    # Note that the presigned URL already contains this in cleartext (we don't want want to rely on it,
    # but it shows that we don't have to consider this information as sensitive).
    multipart_upload_id: str
    part_number: int
    destination_object_key: ObjectStorageKey


@dataclass
class AzureBlobStoragePartExportRequest(SerdeMixin):
    # Source
    source_object_key: ObjectStorageKey  # S3 object in the Terality export bucket
    range_start_byte: int
    range_end_byte: int

    # Destination
    destination_object_key: ObjectStorageKey  # (container, name) of an Azure Blob in storage_account_name
    storage_account_name: str
    shared_access_signature: str


@dataclass
class AwsS3PartsExport(SerdeMixin):
    parts: List[AwsS3ObjectPartExportRequest]


@dataclass
class AzureBlobStoragePartsExport(SerdeMixin):
    parts: List[AzureBlobStoragePartExportRequest]


@dataclass
class ExportToCloudRequest(SerdeMixin):
    service: StorageService
    export_request: Union[AwsS3PartsExport, AzureBlobStoragePartsExport]


@dataclass
class UploadedS3PartInfo(SerdeMixin):
    destination_object_key: ObjectStorageKey
    etag: str
    part_number: int
    multipart_upload_id: str


@dataclass
class ExportToS3Response(SerdeMixin):
    uploaded_parts: List[UploadedS3PartInfo]


@dataclass
class ExportToAzureBlobStorageResponse(SerdeMixin):
    pass


@dataclass
class CreateUserRequest(SerdeMixin):
    email: str
    user_accepted_privacy_policy: bool
    user_accepted_terms_of_service: bool
    full_name: Optional[str]
    company: Optional[str]


@dataclass
class CreateUserResponse(SerdeMixin):
    api_key: str


@dataclass
class ErrorResponse(SerdeMixin):
    message: str


@dataclass
class ClientErrorContext(SerdeMixin):
    code: str
    filename: str
    exception_str: Optional[str]


@dataclass
class CreateAnonymousUserResponse(SerdeMixin):
    user_id: str
    api_key: str
    expires_at: dt.datetime
