from datetime import UTC, datetime, timedelta
from typing import Protocol

from azure.storage.blob import BlobSasPermissions, generate_blob_sas
from azure.storage.blob.aio import BlobServiceClient


class AzureStorage:
    def __init__(self, connection_string: str | None):
        if connection_string:
            self.client = BlobServiceClient.from_connection_string(connection_string)

    def _generate_sas_token(self, filename: str, bucket: str, minute_duration: int):
        expiry_time = datetime.now(UTC) + timedelta(minutes=minute_duration)

        token = generate_blob_sas(
            account_name=self.client.account_name,  # type: ignore
            container_name=bucket,
            blob_name=filename,
            account_key=self.client.credential.account_key,
            permission=BlobSasPermissions(read=True),
            expiry=expiry_time,
        )

        return token

    async def get_file_url(
        self, filename: str, bucket: str, minute_duration: int, image: bool
    ) -> str | None:
        if not hasattr(self, "client"):
            return None

        blob_client = self.client.get_blob_client(container=bucket, blob=filename)

        exists = await blob_client.exists()

        if exists:
            sas_token = self._generate_sas_token(filename, bucket, minute_duration)
            return f"{blob_client.url}?{sas_token}"
        else:
            return None

    async def get_blob_bytes(self, bucket: str, filename: str):
        if not hasattr(self, "client"):
            raise ValueError("No connection string provided to AzureStorage object.")
        blob_client = self.client.get_blob_client(container=bucket, blob=filename)
        return (await blob_client.download_blob()).readall()

    class Config(Protocol):
        azure_blob_connection_string: str

    @classmethod
    def from_config(cls, c: Config):
        return cls(c.azure_blob_connection_string)
