mtsd/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
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


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 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()