shanxia2.0/backend/app/core/storage.py

"""文件存储工具模块

支持多种存储方式：
1. 本地存储 (local) - 默认
2. 七牛云存储 (qiniu) - 推荐生产环境使用
3. 阿里云OSS (oss) - 可选

使用方式:
    from app.core.storage import storage

    # 上传文件
    url = storage.upload_file(file_content, "test.jpg", "images")

    # 上传 Base64 图片
    url = storage.upload_base64(base64_data, "avatar.png", "avatar")

    # 删除文件
    storage.delete_file("images/test.jpg")
"""

import os
import uuid
import base64
from io import BytesIO
from datetime import datetime, timedelta
from urllib.parse import quote
from typing import Optional, BinaryIO
from pathlib import Path

from app.core.config import settings


# HEIC 格式的 MIME 类型
HEIC_TYPES = {"image/heic", "image/heif"}


def convert_heic_to_jpeg(content: bytes) -> bytes:
    """将 HEIC/HEIF 格式转换为 JPEG 格式"""
    try:
        import pillow_heif
        from PIL import Image

        heif_file = pillow_heif.open(BytesIO(content))
        image = heif_file.to_pillow()
        output = BytesIO()
        image.save(output, format="JPEG", quality=85)
        return output.getvalue()
    except ImportError:
        raise Exception("请安装 pillow-heif: pip install pillow-heif")


def process_image_content(content: bytes, filename: str) -> tuple[bytes, str]:
    """处理图片内容，如果是 HEIC 格式则转换为 JPEG"""
    import imghdr
    img_type = imghdr.what(None, h=content[:32])
    if img_type in ("heic", "heif"):
        return convert_heic_to_jpeg(content), filename.rsplit(".", 1)[0] + ".jpg"
    return content, filename


class BaseStorage:
    """存储基类"""

    def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
        """上传文件"""
        raise NotImplementedError

    def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
        """上传 Base64 编码的图片"""
        raise NotImplementedError

    def delete_file(self, file_path: str) -> bool:
        """删除文件"""
        raise NotImplementedError

    def get_url(self, file_path: str) -> str:
        """获取文件访问URL"""
        raise NotImplementedError

    def presigned_download_url(
        self, file_path: str, download_name: str, expires_seconds: int = 3600
    ) -> str:
        """带下载文件名的访问 URL；默认实现为普通访问地址。"""
        return self.get_url(file_path)

    def presigned_preview_url(self, file_path: str, expires_seconds: int = 3600) -> str:
        """浏览器内联预览 URL；默认实现为普通访问地址。"""
        return self.get_url(file_path)


class LocalStorage(BaseStorage):
    """本地存储"""

    def __init__(self):
        self.base_url = settings.LOCAL_STORAGE_URL.rstrip("/")
        self.upload_dir = settings.LOCAL_UPLOAD_DIR

    def _ensure_dir(self, folder: str) -> Path:
        """确保目录存在"""
        path = Path(self.upload_dir) / folder / datetime.now().strftime("%Y%m")
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _generate_filename(self, filename: str) -> str:
        """生成唯一文件名"""
        ext = os.path.splitext(filename)[1] or ""
        unique_name = f"{uuid.uuid4().hex}{ext}"
        return unique_name

    def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
        """上传文件到本地"""
        # 处理 HEIC 格式图片
        file_content, filename = process_image_content(file_content, filename)

        folder_path = self._ensure_dir(folder)
        unique_filename = self._generate_filename(filename)
        file_path = folder_path / unique_filename

        with open(file_path, "wb") as f:
            f.write(file_content)

        relative_path = f"{folder}/{datetime.now().strftime('%Y%m')}/{unique_filename}" if folder else f"{datetime.now().strftime('%Y%m')}/{unique_filename}"
        return self.get_url(relative_path)

    def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
        """上传 Base64 图片"""
        # 去除 data:image/png;base64, 等前缀
        if "," in base64_data:
            base64_data = base64_data.split(",")[1]

        file_content = base64.b64decode(base64_data)
        return self.upload_file(file_content, filename, folder)

    def delete_file(self, file_path: str) -> bool:
        """删除本地文件"""
        try:
            # 移除 URL 前缀获取相对路径
            if file_path.startswith(self.base_url):
                relative_path = file_path.replace(self.base_url, "").lstrip("/")
            else:
                relative_path = file_path.lstrip("/")

            file_full_path = Path(self.upload_dir) / relative_path
            if file_full_path.exists():
                file_full_path.unlink()
                return True
            return False
        except Exception:
            return False

    def get_url(self, file_path: str) -> str:
        """获取文件访问URL"""
        return f"{self.base_url}/{file_path}"


class QiniuStorage(BaseStorage):
    """七牛云存储"""

    def __init__(self):
        try:
            from qiniu import Auth, put_data, delete
            self.qiniu = Auth(settings.QINIU_ACCESS_KEY, settings.QINIU_SECRET_KEY)
            self.bucket = settings.QINIU_BUCKET
            self.base_url = settings.QINIU_DOMAIN.rstrip("/")
        except ImportError:
            raise ImportError("请安装七牛云 SDK: pip install qiniu")

    def _generate_filename(self, filename: str) -> str:
        """生成唯一文件名"""
        ext = os.path.splitext(filename)[1] or ""
        unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
        return unique_name

    def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
        """上传文件到七牛云"""
        from qiniu import put_data

        # 处理 HEIC 格式图片
        file_content, filename = process_image_content(file_content, filename)

        key = self._generate_filename(filename)
        token = self.qiniu.upload_token(self.bucket, key, 3600)

        ret, info = put_data(token, key, file_content)
        if ret is not None:
            return self.get_url(ret["key"])
        raise Exception(f"七牛云上传失败: {info}")

    def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
        """上传 Base64 图片到七牛云"""
        if "," in base64_data:
            base64_data = base64_data.split(",")[1]

        file_content = base64.b64decode(base64_data)
        return self.upload_file(file_content, filename, folder)

    def delete_file(self, file_path: str) -> bool:
        """删除七牛云文件"""
        from qiniu import delete

        try:
            key = file_path.replace(f"{self.base_url}/", "")
            ret, info = delete(self.bucket, key)
            return ret is not None
        except Exception:
            return False

    def get_url(self, file_path: str) -> str:
        """获取文件访问URL"""
        return f"{self.base_url}/{file_path}"


class OSSStorage(BaseStorage):
    """阿里云OSS存储"""

    def __init__(self):
        try:
            import oss2
            self.auth = oss2.Auth(settings.OSS_ACCESS_KEY_ID, settings.OSS_ACCESS_KEY_SECRET)
            self.bucket_name = settings.OSS_BUCKET
            self.bucket = oss2.Bucket(self.auth, settings.OSS_ENDPOINT, self.bucket_name)
            self.base_url = settings.OSS_DOMAIN.rstrip("/")
        except ImportError:
            raise ImportError("请安装阿里云 OSS SDK: pip install oss2")

    def _generate_filename(self, filename: str) -> str:
        """生成唯一文件名"""
        ext = os.path.splitext(filename)[1] or ""
        unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
        return unique_name

    def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
        """上传文件到OSS"""
        # 处理 HEIC 格式图片
        file_content, filename = process_image_content(file_content, filename)

        key = self._generate_filename(filename)
        result = self.bucket.put_object(key, file_content)
        if result.status == 200:
            return self.get_url(key)
        raise Exception(f"OSS上传失败: {result.status}")

    def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
        """上传 Base64 图片到OSS"""
        if "," in base64_data:
            base64_data = base64_data.split(",")[1]

        file_content = base64.b64decode(base64_data)
        return self.upload_file(file_content, filename, folder)

    def delete_file(self, file_path: str) -> bool:
        """删除OSS文件"""
        try:
            key = file_path.replace(f"{self.base_url}/", "")
            result = self.bucket.delete_object(key)
            return result.status == 204
        except Exception:
            return False

    def get_url(self, file_path: str) -> str:
        """获取文件访问URL"""
        return f"{self.base_url}/{file_path}"


class MinIOStorage(BaseStorage):
    """MinIO存储"""

    def __init__(self):
        try:
            from minio import Minio
            self.client = Minio(
                settings.MINIO_ENDPOINT,
                access_key=settings.MINIO_ACCESS_KEY,
                secret_key=settings.MINIO_SECRET_KEY,
                secure=settings.MINIO_SECURE,
            )
            self.bucket_name = settings.MINIO_BUCKET
            if settings.MINIO_BASE_URL:
                self.base_url = settings.MINIO_BASE_URL.rstrip("/")
            else:
                scheme = "https" if settings.MINIO_SECURE else "http"
                self.base_url = f"{scheme}://{settings.MINIO_ENDPOINT}/{self.bucket_name}"
        except ImportError:
            raise ImportError("请安装 MinIO SDK: pip install minio")

    def _generate_filename(self, filename: str) -> str:
        """生成唯一文件名"""
        ext = os.path.splitext(filename)[1] or ""
        unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
        return unique_name

    def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
        """上传文件到MinIO"""
        # 处理 HEIC 格式图片
        file_content, filename = process_image_content(file_content, filename)

        key = self._generate_filename(filename)
        if folder:
            key = f"{folder.strip('/')}/{key}"
        data = BytesIO(file_content)
        self.client.put_object(
            bucket_name=self.bucket_name,
            object_name=key,
            data=data,
            length=len(file_content),
        )
        return self.get_url(key)

    def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
        """上传 Base64 图片到MinIO"""
        if "," in base64_data:
            base64_data = base64_data.split(",")[1]

        file_content = base64.b64decode(base64_data)
        return self.upload_file(file_content, filename, folder)

    def delete_file(self, file_path: str) -> bool:
        """删除MinIO文件"""
        try:
            key = file_path.replace(f"{self.base_url}/", "")
            self.client.remove_object(self.bucket_name, key)
            return True
        except Exception:
            return False

    def get_url(self, file_path: str) -> str:
        """获取文件访问URL"""
        return f"{self.base_url}/{file_path}"

    def _extract_object_key(self, file_path: str) -> str:
        u = (file_path or "").strip()
        if not u:
            return ""
        if "?" in u:
            u = u.split("?", 1)[0]
        prefix = f"{self.base_url}/"
        if u.startswith(prefix):
            return u[len(prefix) :]
        if u.startswith("http://") or u.startswith("https://"):
            marker = f"/{self.bucket_name}/"
            idx = u.find(marker)
            if idx >= 0:
                return u[idx + len(marker) :]
        return u.lstrip("/")

    def presigned_download_url(
        self, file_path: str, download_name: str, expires_seconds: int = 3600
    ) -> str:
        """MinIO 预签名下载 URL，携带原始文件名。"""
        key = self._extract_object_key(file_path)
        if not key:
            raise ValueError("无效的文件路径")
        name = (download_name or "").strip() or key.rsplit("/", 1)[-1] or "download"
        ascii_name = name.encode("ascii", "ignore").decode() or "download"
        encoded = quote(name, safe="")
        disposition = f'attachment; filename="{ascii_name}"; filename*=UTF-8\'\'{encoded}'
        return self.client.presigned_get_object(
            bucket_name=self.bucket_name,
            object_name=key,
            expires=timedelta(seconds=expires_seconds),
            response_headers={"response-content-disposition": disposition},
        )

    def presigned_preview_url(self, file_path: str, expires_seconds: int = 3600) -> str:
        """MinIO 预签名内联预览 URL（PDF 等 iframe 场景）。"""
        key = self._extract_object_key(file_path)
        if not key:
            raise ValueError("无效的文件路径")
        response_headers: dict[str, str] = {"response-content-disposition": "inline"}
        if key.lower().endswith(".pdf"):
            response_headers["response-content-type"] = "application/pdf"
        return self.client.presigned_get_object(
            bucket_name=self.bucket_name,
            object_name=key,
            expires=timedelta(seconds=expires_seconds),
            response_headers=response_headers,
        )


def get_storage() -> BaseStorage:
    """获取存储实例"""
    storage_type = settings.STORAGE_TYPE.lower()

    if storage_type == "qiniu":
        return QiniuStorage()
    elif storage_type == "oss":
        return OSSStorage()
    elif storage_type == "minio":
        return MinIOStorage()
    else:
        return LocalStorage()


# 全局存储实例
storage = get_storage()