zkxw/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 List, 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)


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 _internal_base_prefixes(self) -> List[str]:
        """SDK 预签名 URL 使用的内网地址前缀（http/https 均覆盖）。"""
        endpoint = (settings.MINIO_ENDPOINT or "").strip()
        if not endpoint:
            return []
        bucket = self.bucket_name
        return [
            f"http://{endpoint}/{bucket}",
            f"https://{endpoint}/{bucket}",
        ]

    def _rewrite_presigned_public_url(self, url: str) -> str:
        """将 SDK 生成的内网 MinIO 预签名 URL 改写为 MINIO_BASE_URL 对外路径。"""
        if not url:
            return url
        public_base = (self.base_url or "").rstrip("/")
        if not public_base:
            return url
        for prefix in self._internal_base_prefixes():
            if url.startswith(prefix):
                return public_base + url[len(prefix) :]
        return url

    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}'
        url = self.client.presigned_get_object(
            bucket_name=self.bucket_name,
            object_name=key,
            expires=timedelta(seconds=expires_seconds),
            response_headers={"response-content-disposition": disposition},
        )
        return self._rewrite_presigned_public_url(url)


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