# src/beautyspot/storage.py

import os
import io
import logging
from pathlib import Path
from abc import ABC, abstractmethod
from typing import Any, TypeAlias, Iterator, Protocol, runtime_checkable
from dataclasses import dataclass
from beautyspot.exceptions import CacheCorruptedError, ValidationError

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    boto3 = None
    ClientError = Exception


ReadableBuffer: TypeAlias = bytes | bytearray | memoryview

# --- Storage Policies ---


@runtime_checkable
class StoragePolicyProtocol(Protocol):
    """
    Protocol to determine if data should be saved as a blob (file/object storage)
    or directly in the database based on the data content (usually size).
    """

    def should_save_as_blob(self, data: bytes) -> bool: ...


@dataclass
class ThresholdStoragePolicy(StoragePolicyProtocol):
    """
    Policy that saves data as a blob if its size exceeds a configured threshold.
    This is the recommended policy for automatic optimization.
    """

    threshold: int

    def should_save_as_blob(self, data: bytes) -> bool:
        return len(data) > self.threshold


@dataclass
class WarningOnlyPolicy(StoragePolicyProtocol):
    """
    Policy for backward compatibility (v2.0 behavior).
    Does not force blob storage, but logs a warning if size exceeds threshold.
    """

    warning_threshold: int
    logger: logging.Logger

    def should_save_as_blob(self, data: bytes) -> bool:
        if len(data) > self.warning_threshold:
            self.logger.warning(
                f"⚠️ Large data detected ({len(data)} bytes). "
                f"Consider using `save_blob=True` or a stricter StoragePolicy."
            )
        return False


@dataclass
class AlwaysBlobPolicy(StoragePolicyProtocol):
    """
    Policy that always saves data as a blob.
    Equivalent to setting `default_save_blob=True`.
    """

    def should_save_as_blob(self, data: bytes) -> bool:
        return True


# --- Blob Storage Implementations ---

class BlobStorageBase(ABC):
    """
    Abstract base class for large object storage (BLOBs).
    """

    @abstractmethod
    def save(self, key: str, data: ReadableBuffer) -> str:
        """
        Persist the data associated with the given key.
        Returns a location identifier.
        """
        pass

    @abstractmethod
    def load(self, location: str) -> bytes:
        """
        Retrieve data from the specified location.
        """
        pass

    @abstractmethod
    def delete(self, location: str) -> None:
        """
        Delete the blob at the specified location.
        Should be idempotent (no error if file missing).
        """
        pass

    @abstractmethod
    def list_keys(self) -> Iterator[str]:
        """
        Yields location identifiers for all stored blobs.
        Used for Garbage Collection.
        MUST yield the same format (path/URI) that is accepted by `delete`.
        """
        pass


class LocalStorage(BlobStorageBase):
    def __init__(self, base_dir: str | Path):
        # Resolve to absolute path explicitly on init
        self.base_dir = Path(base_dir).resolve()
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def _validate_key(self, key: str):
        # Prevent Path Traversal
        if ".." in key or "/" in key or "\\" in key:
            raise ValidationError(
                f"Invalid key: '{key}'. Keys must not contain path separators."
            )

    def save(self, key: str, data: ReadableBuffer) -> str:
        self._validate_key(key)
        filename = f"{key}.bin"
        filepath = self.base_dir / filename

        # Atomic write: flush + fsync ensures data reaches disk before rename,
        # so a crash between write and rename never leaves a corrupt file.
        temp_path = self.base_dir / f"{filename}.tmp"
        with open(temp_path, "wb") as f:
            f.write(data)
            f.flush()
            os.fsync(f.fileno())

        temp_path.replace(filepath)

        # [CHANGED] Return filename (relative path) instead of absolute path
        # This enables portability of the DB/Blob set across different environments.
        return filename

    def load(self, location: str) -> bytes:
        # [CHANGED] Resolve location relative to base_dir.
        # Note: If 'location' is an absolute path (legacy data), pathlib behavior
        # (base / abs) returns abs, so backward compatibility on the same machine is preserved.
        full_path = (self.base_dir / location).resolve()

        # Security check: Ensure the path is strictly within the base_dir
        if not full_path.is_relative_to(self.base_dir):
            raise ValueError(
                f"Access denied: {location} resolves to {full_path}, which is outside {self.base_dir}"
            )

        if not full_path.exists():
            raise CacheCorruptedError(f"Local blob lost: {full_path}")

        try:
            with open(full_path, "rb") as f:
                return f.read()
        except OSError as e:
            raise CacheCorruptedError(f"Failed to read blob: {e}")

    def delete(self, location: str) -> None:
        """
        Delete the file at the given location.
        
        Note:
            For performance reasons, this method does not synchronously remove 
            empty parent directories. Directory cleanup is deferred to the 
            asynchronous maintenance task (`prune_empty_dirs` / `beautyspot gc`).
        """
        full_path = (self.base_dir / location).resolve()

        if not full_path.is_relative_to(self.base_dir):
            return

        try:
            os.remove(full_path)
        except FileNotFoundError:
            pass

    def list_keys(self) -> Iterator[str]:
        """
        Yields relative paths of all .bin files, including subdirectories.
        Example: 'hash.bin' or 'subdir/hash.bin'
        """
        if not self.base_dir.exists():
            return

        # [変更] rglob で再帰的に探索
        for entry in self.base_dir.rglob("*.bin"):
            if entry.is_file():
                # [変更] base_dir からの相対パスを返す
                yield str(entry.relative_to(self.base_dir))

    def prune_empty_dirs(self) -> int:
        """
        Recursively remove empty directories under base_dir (including base_dir itself).
        Also removes directories containing only system generated files (.DS_Store, etc).
        Returns the count of removed directories.
        """
        if not self.base_dir.exists():
            return 0

        IGNORED_FILES = {".DS_Store", "Thumbs.db", "desktop.ini"}
        removed_count = 0

        # os.walk(topdown=False) で深い階層から順に処理
        for root, dirs, files in os.walk(self.base_dir, topdown=False):
            path = Path(root)

            existing_files = set(files)

            # 無視リスト以外のファイルがある場合 -> 削除不可
            if existing_files - IGNORED_FILES:
                continue

            # 無視リストにあるファイルしか残っていない場合、それらを消して空にする
            for f in existing_files:
                try:
                    (path / f).unlink()
                except OSError:
                    pass

            # ディレクトリ削除を試みる
            try:
                path.rmdir()
                removed_count += 1
            except OSError:
                pass

        return removed_count


class S3Storage(BlobStorageBase):
    def __init__(
        self,
        s3_uri: str,
        s3_opts: dict[str, Any] | None = None,
    ):
        if not boto3:
            raise ImportError("Run `pip install beautyspot[s3]` to use S3 storage.")

        parts = s3_uri.replace("s3://", "").split("/", 1)
        self.bucket_name = parts[0]
        self.prefix = parts[1].rstrip("/") if len(parts) > 1 else "blobs"

        opts = s3_opts or {}
        self.s3 = boto3.client("s3", **opts)

    @staticmethod
    def _parse_s3_uri(location: str) -> tuple[str, str]:
        """Parse an s3:// URI into (bucket, key). Raises ValueError for invalid URIs."""
        if not location.startswith("s3://"):
            raise ValidationError(f"Expected an s3:// URI, got: {location!r}")
        path = location[len("s3://"):]
        parts = path.split("/", 1)
        if len(parts) != 2 or not parts[0] or not parts[1]:
            raise ValidationError(
                f"Invalid S3 URI (expected s3://bucket/key): {location!r}"
            )
        return parts[0], parts[1]

    def save(self, key: str, data: ReadableBuffer) -> str:
        s3_key = f"{self.prefix}/{key}.bin"
        buffer = io.BytesIO(data)
        self.s3.put_object(Bucket=self.bucket_name, Key=s3_key, Body=buffer)
        return f"s3://{self.bucket_name}/{s3_key}"

    def load(self, location: str) -> bytes:
        bucket, key = self._parse_s3_uri(location)
        try:
            resp = self.s3.get_object(Bucket=bucket, Key=key)
            return resp["Body"].read()
        except ClientError as e:
            raise CacheCorruptedError(f"S3 blob lost: {location}") from e

    def delete(self, location: str) -> None:
        bucket, key = self._parse_s3_uri(location)
        try:
            self.s3.delete_object(Bucket=bucket, Key=key)
        except ClientError:
            pass

    def list_keys(self) -> Iterator[str]:
        """Yields s3:// URIs for all objects in the prefix."""
        paginator = self.s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self.bucket_name, Prefix=self.prefix):
            for obj in page.get("Contents", []):
                yield f"s3://{self.bucket_name}/{obj['Key']}"


def create_storage(path: str, options: dict | None = None) -> BlobStorageBase:
    if path.startswith("s3://"):
        return S3Storage(path, options)

    return LocalStorage(path)
