"""Pydantic models for email composer attachments."""

import base64
from pydantic import Field, BaseModel

from common.src import aws
from common.src.aws.utils import get_s3_bucket_key_from_uri
from common.src.typings import Base64String, S3URI

from ... import config


class BaseAttachment(BaseModel):
    """Base attachment model. In the future it can be subclassed
    to e.g. S3Attachment.

    ALL SUBCLASSES MUST IMPLEMENT the `attachment_data` property, and may
    also override the `is_b64_string` property.
    """

    file_name: str = Field(
        ...,
        description="Name of the attachment file, including extension",
        example="my_file.xlsx",
    )
    mime_type: str = Field(
        ...,
        description="MIME type of the attachment file",
        example="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    )

    @property
    def is_b64_string(self) -> bool:
        """Whether value of `attachment` is already a Base64 string."""
        return False  # Default for subclasses


class Base64Attachment(BaseAttachment):
    """Attachment where the provided data is a valid Base64 encoded string."""

    attachment: Base64String = Field(
        ..., description="Base64-encoded attachment payload"
    )

    @property
    def is_b64_string(self) -> bool:
        """Whether value of `attachment` is a Base64 string."""
        return True

    @property
    def attachment_data(self) -> bytes:
        """Decode attachment from Base64 and return as bytes."""
        return base64.b64decode(self.attachment)


class S3Attachment(BaseAttachment):
    """Attachment where the provided data is a valid S3 URI.

    The attachment is downloaded from S3 and returned as bytes.
    """

    attachment: S3URI = Field(
        ...,
        description="Full S3 path to the attachment file",
        example="s3://bucket/path/to/file.xlsx",
    )

    @property
    def attachment_data(self) -> bytes:
        """Download attachment from S3 and return as bytes."""
        session = aws.Session()
        s3_client = session.client(aws.AWSServices.S3, region_name=config.AWS_REGION)
        s3_bucket, s3_blob = get_s3_bucket_key_from_uri(self.attachment)
        response = s3_client.get_object(Bucket=s3_bucket, Key=s3_blob)
        file_content = response["Body"].read()
        return file_content