"""
Async S3/MinIO client. Abstracts the details so application code never
touches boto3 directly.
"""
from __future__ import annotations

import hashlib
import io
import logging
import uuid
from datetime import datetime, timezone
from typing import AsyncIterator, BinaryIO

import aioboto3
from botocore.config import Config as BotoConfig
from botocore.exceptions import ClientError
from pydantic_settings import BaseSettings

logger = logging.getLogger(__name__)


class S3Settings(BaseSettings):
    s3_endpoint_url: str
    """Host the API uses to reach S3 (e.g. http://minio:9000 inside Docker)."""
    s3_public_endpoint_url: str | None = None
    """If set, presigned GET/PUT URLs use this host so browsers can open them (e.g. http://localhost:19000)."""
    s3_region: str = "us-east-1"
    s3_access_key: str
    s3_secret_key: str
    s3_bucket_attachments: str
    s3_bucket_exports: str
    s3_bucket_quarantine: str
    s3_bucket_templates: str = "initiative-templates"
    """Admin-managed .docx templates (server-side only; no browser CORS needed)."""
    s3_bucket_imagehub_blobs: str = "imagehub-blobs"
    """ImageHub content-addressed blob store (imaging dataset files; deduped by sha256)."""
    s3_signed_url_ttl: int = 900  # 15 minutes
    max_upload_size_mb: int = 50
    max_blob_size_mb: int = 2048  # ImageHub imaging blobs are large (DICOM series, NIfTI volumes)

    class Config:
        env_file = ".env"


settings = S3Settings()


# Allowed MIME types — validated server-side, never trust the client
ALLOWED_MIME_TYPES = {
    "application/pdf",
    "image/png",
    "image/jpeg",
    "image/webp",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",  # docx
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",         # xlsx
    "application/vnd.openxmlformats-officedocument.presentationml.presentation", # pptx
    "application/msword",
    "application/vnd.ms-excel",
    "text/plain",
}


class S3Storage:
    """
    Async context-manager friendly wrapper around MinIO/S3.

    Design choices:
    - Keys are content-addressed: {app_id}/{yyyy}/{mm}/{uuid}-{safe_filename}
      This avoids collisions, makes listing by app cheap, and sharding by
      month keeps directory-style prefixes from getting too deep.
    - Never streams a whole file into memory — chunked upload/download
      for large attachments.
    - All errors are caught and re-raised as domain exceptions.
    """

    def __init__(self, settings: S3Settings = settings):
        self._settings = settings
        self._session = aioboto3.Session(
            aws_access_key_id=settings.s3_access_key,
            aws_secret_access_key=settings.s3_secret_key,
            region_name=settings.s3_region,
        )

    def _client(self):
        """Return a client as an async context manager."""
        return self._session.client(
            "s3",
            endpoint_url=self._settings.s3_endpoint_url,
            # Path-style addressing works with both MinIO and S3
            config=BotoConfig(signature_version="s3v4"),
        )

    def _client_presign(self):
        """Presigned URLs must use a host reachable from the user's browser (often not minio:9000)."""
        ep = self._settings.s3_public_endpoint_url or self._settings.s3_endpoint_url
        return self._session.client(
            "s3",
            endpoint_url=ep,
            config=BotoConfig(signature_version="s3v4"),
        )

    async def ensure_buckets_exist(self) -> None:
        """Create configured buckets if they do not exist (local MinIO / first boot)."""
        names = (
            self._settings.s3_bucket_attachments,
            self._settings.s3_bucket_exports,
            self._settings.s3_bucket_quarantine,
            self._settings.s3_bucket_templates,
            self._settings.s3_bucket_imagehub_blobs,
        )
        for name in names:
            async with self._client() as s3:
                try:
                    await s3.create_bucket(Bucket=name)
                except ClientError as exc:
                    code = (exc.response or {}).get("Error", {}).get("Code", "")
                    if code in (
                        "BucketAlreadyOwnedByYou",
                        "BucketAlreadyExists",
                    ):
                        continue
                    raise
            logger.info("S3 bucket ready: %s", name)

    # ------------------------------------------------------------------ #
    # Key construction
    # ------------------------------------------------------------------ #
    @staticmethod
    def build_key(application_id: int, filename: str) -> str:
        """
        Build a unique, safe object key for an attachment.

        Example: 42/2025/10/7c9e8a1b-4d5f-flowchart.pdf
        """
        now = datetime.now(tz=timezone.utc)
        safe = _sanitize_filename(filename)
        unique = uuid.uuid4().hex[:16]
        return f"{application_id}/{now:%Y}/{now:%m}/{unique}-{safe}"

    @staticmethod
    def build_key_for_initiative(initiative_id: uuid.UUID, filename: str) -> str:
        """
        Same as build_key but namespaces by initiative UUID (initiative application drafts).
        """
        now = datetime.now(tz=timezone.utc)
        safe = _sanitize_filename(filename)
        unique = uuid.uuid4().hex[:16]
        iid = str(initiative_id).replace("-", "")
        return f"initiatives/{iid}/{now:%Y}/{now:%m}/{unique}-{safe}"

    # ------------------------------------------------------------------ #
    # Upload
    # ------------------------------------------------------------------ #
    async def upload(
        self,
        bucket: str,
        key: str,
        fileobj: BinaryIO,
        mime_type: str,
        metadata: dict[str, str] | None = None,
    ) -> dict:
        """
        Upload a file-like object to MinIO. Validates MIME type and size,
        computes SHA-256 for integrity, returns result metadata.
        """
        if mime_type not in ALLOWED_MIME_TYPES:
            raise ValueError(f"MIME type not allowed: {mime_type}")

        # Compute size + hash while buffering
        data = fileobj.read()
        size = len(data)
        max_bytes = self._settings.max_upload_size_mb * 1024 * 1024
        if size > max_bytes:
            raise ValueError(
                f"File too large: {size} bytes > {max_bytes} bytes limit"
            )
        sha256 = hashlib.sha256(data).hexdigest()

        meta = {"sha256": sha256, **(metadata or {})}

        async with self._client() as s3:
            try:
                await s3.put_object(
                    Bucket=bucket,
                    Key=key,
                    Body=io.BytesIO(data),
                    ContentType=mime_type,
                    Metadata=meta,
                    # Omit ServerSideEncryption: default MinIO returns NotImplemented if KMS/SSE
                    # is not configured. Use bucket policies or a compatible backend for SSE.
                )
            except ClientError as exc:
                logger.exception("S3 upload failed: bucket=%s key=%s", bucket, key)
                raise StorageError(f"Upload failed: {exc}") from exc

        logger.info(
            "Uploaded s3://%s/%s size=%d sha256=%s", bucket, key, size, sha256[:12]
        )
        return {"bucket": bucket, "key": key, "size": size, "sha256": sha256}

    # ------------------------------------------------------------------ #
    # Download (streamed)
    # ------------------------------------------------------------------ #
    async def download_stream(
        self, bucket: str, key: str, chunk_size: int = 64 * 1024
    ) -> AsyncIterator[bytes]:
        """
        Stream the object body in chunks. Use for FastAPI StreamingResponse.
        """
        async with self._client() as s3:
            try:
                obj = await s3.get_object(Bucket=bucket, Key=key)
            except ClientError as exc:
                if exc.response["Error"]["Code"] == "NoSuchKey":
                    raise FileNotFoundError(f"Object not found: {key}") from exc
                raise StorageError(f"Download failed: {exc}") from exc

            # aiobotocore: `async with body as stream` yields aiohttp ClientResponse;
            # `stream.read(n)` fails (read() takes no size). Keep the Body wrapper and
            # read from `body` inside `async with body:` so chunk sizes work. See aiobotocore#1156.
            body = obj["Body"]
            async with body:
                while True:
                    chunk = await body.read(chunk_size)
                    if not chunk:
                        break
                    yield chunk

    # ------------------------------------------------------------------ #
    # Signed URLs
    # ------------------------------------------------------------------ #
    async def get_download_url(
        self,
        bucket: str,
        key: str,
        ttl: int | None = None,
        filename: str | None = None,
        *,
        inline: bool = False,
        response_content_type: str | None = None,
    ) -> str:
        """
        Generate a pre-signed URL so the browser can GET directly from MinIO.

        - ``inline=False`` (default): Content-Disposition attachment — save-as / download.
        - ``inline=True``: Content-Disposition inline — PDF/image viewers and iframes.
        - ``response_content_type``: optional override (e.g. application/pdf) for clients
          that rely on the response header when the stored object metadata is wrong.
        """
        ttl = ttl or self._settings.s3_signed_url_ttl
        params: dict = {"Bucket": bucket, "Key": key}
        if filename:
            safe = _sanitize_filename(filename)
            disp = "inline" if inline else "attachment"
            params["ResponseContentDisposition"] = f'{disp}; filename="{safe}"'
        elif inline:
            params["ResponseContentDisposition"] = "inline"
        if response_content_type:
            params["ResponseContentType"] = response_content_type
        async with self._client_presign() as s3:
            return await s3.generate_presigned_url(
                "get_object", Params=params, ExpiresIn=ttl
            )

    async def get_upload_url(
        self, bucket: str, key: str, mime_type: str, ttl: int | None = None
    ) -> dict:
        """
        Generate a pre-signed URL for direct browser → MinIO upload.
        Returns URL + required headers. Use this for large files to bypass
        the FastAPI request-body limit.
        """
        ttl = ttl or self._settings.s3_signed_url_ttl
        async with self._client_presign() as s3:
            url = await s3.generate_presigned_url(
                "put_object",
                Params={
                    "Bucket": bucket,
                    "Key": key,
                    "ContentType": mime_type,
                },
                ExpiresIn=ttl,
            )
        return {"url": url, "headers": {"Content-Type": mime_type}}

    # ------------------------------------------------------------------ #
    # Delete / copy
    # ------------------------------------------------------------------ #
    async def delete(self, bucket: str, key: str) -> None:
        """
        Delete an object. With versioning enabled, this creates a delete
        marker — the previous version is still recoverable until lifecycle
        rules remove it.
        """
        async with self._client() as s3:
            await s3.delete_object(Bucket=bucket, Key=key)
        logger.info("Deleted s3://%s/%s", bucket, key)

    async def move(
        self, src_bucket: str, src_key: str, dst_bucket: str, dst_key: str
    ) -> None:
        """Atomic move: copy then delete. Used when quarantining suspicious files."""
        async with self._client() as s3:
            await s3.copy_object(
                Bucket=dst_bucket,
                Key=dst_key,
                CopySource={"Bucket": src_bucket, "Key": src_key},
            )
            await s3.delete_object(Bucket=src_bucket, Key=src_key)

    async def head(self, bucket: str, key: str) -> dict:
        """Fetch object metadata without downloading it."""
        async with self._client() as s3:
            try:
                return await s3.head_object(Bucket=bucket, Key=key)
            except ClientError as exc:
                if exc.response["Error"]["Code"] == "404":
                    raise FileNotFoundError(f"Object not found: {key}") from exc
                raise

    # ------------------------------------------------------------------ #
    # Content-addressed blobs (ImageHub) — dedup by sha256
    # ------------------------------------------------------------------ #
    @staticmethod
    def build_blob_key(sha256: str) -> str:
        """Content-addressed key ``blobs/<aa>/<bb>/<sha256>`` (sharded by hash prefix)."""
        h = (sha256 or "").lower()
        aa = h[0:2] if len(h) >= 2 else "00"
        bb = h[2:4] if len(h) >= 4 else "00"
        return f"blobs/{aa}/{bb}/{h}"

    async def blob_exists(self, bucket: str, key: str) -> bool:
        """True if an object already exists at key (used for content-addressed dedup)."""
        async with self._client() as s3:
            try:
                await s3.head_object(Bucket=bucket, Key=key)
                return True
            except ClientError as exc:
                code = (exc.response or {}).get("Error", {}).get("Code", "")
                if code in ("404", "NoSuchKey", "NotFound"):
                    return False
                raise

    async def put_blob(self, data: bytes, media_type: str | None = None) -> dict:
        """
        Store raw bytes as a content-addressed, globally deduped blob.

        Hashes the bytes (sha256) and PUTs to ``blobs/<aa>/<bb>/<hash>`` in the ImageHub
        bucket only if not already present (``deduped`` tells the caller which happened).
        Accepts any ``media_type`` — imaging MIME is unreliable (DICOM is often
        ``application/octet-stream``) — but enforces ``max_blob_size_mb``.

        NOTE: like ``upload()``, this buffers the whole blob in memory; resumable/
        multipart upload for very large series is a later milestone.
        """
        size = len(data)
        max_bytes = self._settings.max_blob_size_mb * 1024 * 1024
        if size > max_bytes:
            raise ValueError(f"Blob too large: {size} bytes > {max_bytes} bytes limit")
        sha256 = hashlib.sha256(data).hexdigest()
        bucket = self._settings.s3_bucket_imagehub_blobs
        key = self.build_blob_key(sha256)
        ctype = media_type or "application/octet-stream"
        if await self.blob_exists(bucket, key):
            return {"sha256": sha256, "size": size, "bucket": bucket, "key": key,
                    "media_type": ctype, "deduped": True}
        async with self._client() as s3:
            try:
                await s3.put_object(
                    Bucket=bucket, Key=key, Body=io.BytesIO(data), ContentType=ctype,
                    Metadata={"sha256": sha256},
                )
            except ClientError as exc:
                logger.exception("S3 put_blob failed: bucket=%s key=%s", bucket, key)
                raise StorageError(f"Blob upload failed: {exc}") from exc
        logger.info("Stored blob s3://%s/%s size=%d sha256=%s", bucket, key, size, sha256[:12])
        return {"sha256": sha256, "size": size, "bucket": bucket, "key": key,
                "media_type": ctype, "deduped": False}


# ---------------------------------------------------------------------- #
# Helpers
# ---------------------------------------------------------------------- #
def _sanitize_filename(name: str) -> str:
    """
    Strip path components and replace unsafe characters. Preserves
    Vietnamese diacritics because MinIO keys are UTF-8.
    """
    import re
    import unicodedata

    # Strip any path separators
    name = name.replace("/", "_").replace("\\", "_")
    # Collapse whitespace
    name = re.sub(r"\s+", "_", name.strip())
    # Remove control chars
    name = "".join(ch for ch in name if unicodedata.category(ch)[0] != "C")
    # Max length 200 chars (S3 key limit is 1024, but be conservative)
    return name[:200] or "file"


class StorageError(Exception):
    """Raised for any storage-layer failure."""


# Module-level singleton
storage = S3Storage()
