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 mimetypes
import os
import uuid
import base64
from io import BytesIO
from datetime import datetime, timedelta
from urllib.parse import quote, urlparse
from typing import Any, Optional, BinaryIO
from pathlib import Path

from app.core.config import settings


def _guess_preview_content_type(object_key: str) -> Optional[str]:
    """按对象键后缀推断内联预览 Content-Type（纠正上传时未写入 MIME 的情况）。"""
    path = (object_key or "").split("?", 1)[0].lower()
    # 常见图片/PDF 显式映射，避免依赖系统 mime 库差异
    explicit = {
        ".pdf": "application/pdf",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".bmp": "image/bmp",
        ".webp": "image/webp",
        ".svg": "image/svg+xml",
    }
    for ext, mime in explicit.items():
        if path.endswith(ext):
            return mime
    guessed, _ = mimetypes.guess_type(path)
    if guessed and guessed.startswith(("image/", "application/pdf", "text/", "video/")):
        return guessed
    return None


# 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}"
            # 预签名给浏览器用：优先按 MINIO_BASE_URL 的主机签发，上传/删除仍走内网 client
            self._presign_client = self._build_presign_client(Minio)
        except ImportError:
            raise ImportError("请安装 MinIO SDK: pip install minio")

    def _build_presign_client(self, minio_cls: Any) -> Any:
        """从 MINIO_BASE_URL 解析公网 endpoint；未配置或与内网相同时复用 self.client。"""
        raw = (settings.MINIO_BASE_URL or "").strip()
        if not raw:
            return self.client
        parsed = urlparse(raw if "://" in raw else f"https://{raw}")
        host = parsed.hostname
        if not host:
            return self.client
        endpoint = f"{host}:{parsed.port}" if parsed.port else host
        secure = (parsed.scheme or "https").lower() != "http"
        if endpoint == settings.MINIO_ENDPOINT and secure == bool(settings.MINIO_SECURE):
            return self.client
        # region 固定：预签名不必经公网反代做 GetBucketLocation（反代常只放行 /bucket/object）
        return minio_cls(
            endpoint,
            access_key=settings.MINIO_ACCESS_KEY,
            secret_key=settings.MINIO_SECRET_KEY,
            secure=secure,
            region="us-east-1",
        )

    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._presign_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 等）。

        上传时 put_object 未写 content_type 时对象多为 application/octet-stream，
        浏览器顶级导航会触发下载；此处通过 response-content-type 覆盖为正确 MIME。
        """
        key = self._extract_object_key(file_path)
        if not key:
            raise ValueError("无效的文件路径")
        response_headers: dict[str, str] = {"response-content-disposition": "inline"}
        content_type = _guess_preview_content_type(key)
        if content_type:
            response_headers["response-content-type"] = content_type
        return self._presign_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()