utils/r2_uploader.py

import os
import asyncio
import boto3
from botocore.exceptions import ClientError
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from PIL import Image
import io
import hashlib
import threading
from datetime import datetime, timedelta
from dotenv import load_dotenv

# 配置缓存
_config_cache = {}
_cache_lock = threading.Lock()
_CACHE_TTL = timedelta(minutes=30)  # 缓存 30 分钟

# 全局S3客户端缓存，按公司编码管理
# 格式: {company_code: {"client": s3_client, "created_at": datetime}}
GLOBAL_S3_CLIENTS = {}
# 客户端过期时间（秒）
CLIENT_EXPIRY_SECONDS = 24 * 3600  # 24小时
# 客户端缓存锁
_client_lock = threading.Lock()


# ===================== 先定义类，再使用！=====================
class R2Config:
    """R2 配置对象"""
    def __init__(self, config_data: Dict):
        self.account_id = config_data.get('r2_account_id', '')
        self.access_key_id = config_data.get('r2_access_key_id', '')
        self.secret_access_key = config_data.get('r2_secret_access_key', '')
        self.bucket_name = config_data.get('r2_bucket_name', 'ymt-images')
        self.public_url = config_data.get('r2_public_url', '')
        self.enabled = config_data.get('r2_enabled', 1) == 1
        self.cached_at = datetime.now()

    def is_expired(self) -> bool:
        """检查缓存是否过期"""
        return datetime.now() - self.cached_at > _CACHE_TTL


def get_r2_config_from_env() -> Optional[R2Config]:
    """从 .env 文件获取 R2 配置"""
    # 尝试加载 .env 文件
    try:
        load_dotenv()
    except:
        pass

    r2_account_id = os.getenv('R2_ACCOUNT_ID', '')
    r2_access_key_id = os.getenv('R2_ACCESS_KEY_ID', '')
    r2_secret_access_key = os.getenv('R2_SECRET_ACCESS_KEY', '')
    r2_bucket_name = os.getenv('R2_BUCKET_NAME', 'yomaton')
    r2_public_url = os.getenv('R2_PUBLIC_URL', '')
    r2_enabled = os.getenv('R2_ENABLED', 'true').lower() == 'true'

    if r2_enabled and r2_account_id and r2_access_key_id and r2_secret_access_key:
        return R2Config({
            'r2_account_id': r2_account_id,
            'r2_access_key_id': r2_access_key_id,
            'r2_secret_access_key': r2_secret_access_key,
            'r2_bucket_name': r2_bucket_name,
            'r2_public_url': r2_public_url,
            'r2_enabled': 1 if r2_enabled else 0
        })

    return None


class R2Uploader:
    """
    Cloudflare R2 上传工具类
    """

    def __init__(self, config: R2Config, company_code: str = "default"):
        self.config = config
        self.company_code = company_code
        self._cleanup_expired_clients()

    def _cleanup_expired_clients(self):
        """清理过期的客户端"""
        with _client_lock:
            expired_companies = []
            for company_code, client_info in GLOBAL_S3_CLIENTS.items():
                if (datetime.utcnow() - client_info["created_at"]).total_seconds() > CLIENT_EXPIRY_SECONDS:
                    expired_companies.append(company_code)

            for company_code in expired_companies:
                if company_code != '0000' and company_code != 'default':
                    del GLOBAL_S3_CLIENTS[company_code]
                    print(f"清理过期的S3客户端: {company_code}")

    def _get_config_from_db(self, company_code: str) -> Optional[R2Config]:
        """
        从数据库获取 R2 配置

        注意：这里需要在有数据库会话的上下文调用
        实际使用时通过 get_r2_uploader_for_company() 传入配置
        """
        return None

    def get_client(self):
        """
        获取 R2 S3 客户端，按公司编码全局复用
        优先级：当前公司编码 -> '0000' -> 'default'
        """
        if not self.config.enabled or not self.config.account_id or not self.config.access_key_id or not self.config.secret_access_key:
            return None

        # 清理过期客户端
        self._cleanup_expired_clients()

        # 检查顺序：当前公司 -> 0000 -> default
        client_codes = [self.company_code, '0000', 'default']

        with _client_lock:
            for code in client_codes:
                if code in GLOBAL_S3_CLIENTS:
                    client_info = GLOBAL_S3_CLIENTS[code]
                    time_diff = (datetime.utcnow() - client_info["created_at"]).total_seconds()
                    if time_diff < CLIENT_EXPIRY_SECONDS:
                        print(f"复用S3客户端: {code} (优先级: {client_codes.index(code) + 1})")
                        # 如果不是当前公司的客户端，为当前公司创建一个引用
                        if code != self.company_code:
                            GLOBAL_S3_CLIENTS[self.company_code] = {
                                "client": client_info["client"],
                                "created_at": datetime.utcnow()
                            }
                            print(f"为当前公司 {self.company_code} 创建客户端引用")
                        return client_info["client"]
                    else:
                        # 客户端已过期，删除
                        del GLOBAL_S3_CLIENTS[code]
                        print(f"S3客户端已过期，删除: {code}")

        # 创建新的客户端
        try:
            client = boto3.client(
                's3',
                endpoint_url=f'https://{self.config.account_id}.r2.cloudflarestorage.com',
                aws_access_key_id=self.config.access_key_id,
                aws_secret_access_key=self.config.secret_access_key
            )

            # 缓存客户端
            with _client_lock:
                GLOBAL_S3_CLIENTS[self.company_code] = {
                    "client": client,
                    "created_at": datetime.utcnow()
                }
                # 同时缓存为 'default'（如果是从环境变量加载的配置）
                if self.company_code == 'default':
                    pass  # 已经缓存为 default
                elif self.company_code == '0000':
                    # 0000 公司的配置也缓存为 default
                    GLOBAL_S3_CLIENTS['default'] = {
                        "client": client,
                        "created_at": datetime.utcnow()
                    }

            print(f"创建新的S3客户端: {self.company_code}")
            return client
        except Exception as e:
            print(f"R2 客户端初始化失败: {e}")
            return None

    def is_available(self) -> bool:
        """检查 R2 是否可用"""
        if not self.config or not self.config.enabled:
            return False
        return self.get_client() is not None

    def generate_r2_key(self, company_code: str, product_number: str, file_type: str, filename: str, sub_type: Optional[str] = None) -> str:
        """
        生成 R2 对象键（路径）

        格式: {company_code}/{product_number}/{file_type}/{sub_type}/{filename}
        """
        key_parts = [company_code, product_number, file_type]
        if sub_type:
            key_parts.append(sub_type)
        key_parts.append(filename)
        return '/'.join(key_parts)

    def generate_public_url(self, r2_key: str) -> str:
        """
        生成 R2 公共访问 URL
        """
        if not self.config or not self.config.public_url:
            return ''
        return f"{self.config.public_url.rstrip('/')}/{r2_key.lstrip('/')}"

    def generate_thumbnail(self, file_data: bytes, max_width: int = 400, max_height: int = 400) -> Tuple[bytes, str]:
        """
        从字节数据生成缩略图

        返回: (缩略图二进制数据, 文件扩展名)
        """
        try:
            img = Image.open(io.BytesIO(file_data))

            # 计算缩放尺寸
            img.thumbnail((max_width, max_height), Image.Resampling.LANCZOS)

            # 保存为字节流
            img_byte_arr = io.BytesIO()

            # 确定保存格式
            ext = '.jpg'  # 默认使用JPEG
            if img.mode in ['RGBA', 'P']:
                img = img.convert('RGB')
                ext = '.jpg'
            elif img.mode == 'LA':
                img = img.convert('RGBA')
                ext = '.png'

            save_format = 'JPEG' if ext in ['.jpg', '.jpeg'] else ext[1:].upper()
            img.save(img_byte_arr, format=save_format, quality=85)
            img_byte_arr.seek(0)

            return img_byte_arr.read(), ext
        except Exception as e:
            print(f"缩略图生成失败: {e}")
            return file_data, '.jpg'

    def generate_md5_key(self, file_data: bytes, filename: str) -> str:
        """
        使用文件MD5生成R2键

        参数:
            file_data: 文件字节数据
            filename: 原始文件名

        返回:
            R2键（MD5命名）
        """
        # 计算文件MD5
        hash_md5 = hashlib.md5()
        hash_md5.update(file_data)
        file_hash = hash_md5.hexdigest()

        # 获取文件扩展名
        ext = Path(filename).suffix.lower()

        # 生成MD5命名的R2键
        return f"{file_hash}{ext}"

    async def upload_file(
        self,
        file_buffer: bytes,
        file_name: str,
        content_type: Optional[str] = None,
        generate_thumbnail: bool = True,
        max_retries: int = 3
    ) -> str:
        """
        上传单个文件到 R2

        参数:
            file_buffer: 文件字节数据
            file_name: 文件名（完整R2键，包含路径）
            content_type: MIME类型
            generate_thumbnail: 是否生成缩略图
            max_retries: 最大重试次数

        返回:
            R2 URL
        """
        if not self.is_available():
            return ''

        s3_client = self.get_client()

        # 直接使用传入的文件名作为R2键
        r2_key = file_name

        # 上传原图（带重试机制）
        for attempt in range(max_retries):
            try:
                s3_client.put_object(
                    Bucket=self.config.bucket_name,
                    Key=r2_key,
                    Body=file_buffer,
                    ContentType=content_type or 'application/octet-stream'
                )
                print(f"R2 上传成功: {r2_key}")
                break
            except ClientError as e:
                error_code = e.response.get('Error', {}).get('Code')
                # 对于可重试的错误，进行重试
                if error_code in ['RequestTimeout', 'ConnectionError', 'ServiceUnavailable'] and attempt < max_retries - 1:
                    print(f"R2 上传失败 (尝试 {attempt + 1}/{max_retries}): {e}")
                    await asyncio.sleep(1 * (attempt + 1))  # 指数退避
                else:
                    print(f"R2 上传失败: {e}")
                    raise

        # 生成并上传缩略图（如果是图片，但这里不处理缩略图，因为upload_thumbnail已被单独调用
        # 缩略图将在media.py中单独处理

        return self.generate_public_url(r2_key)

    async def batch_upload_files(
        self,
        files: List[Dict[str, any]],
        max_concurrency: int = 10  # 增加并发数
    ) -> Dict[str, str]:
        """
        批量上传多个文件到 R2

        参数:
            files: 文件列表，每个文件包含 {'file_buffer': bytes, 'file_name': str, 'content_type': Optional[str]}
            max_concurrency: 最大并发数

        返回:
            字典，键为文件名，值为R2 URL
        """
        if not self.is_available():
            return {}

        if not files:
            return {}

        # 获取S3客户端，避免每次上传都创建新连接
        s3_client = self.get_client()
        if not s3_client:
            return {}

        # 限制并发数
        semaphore = asyncio.Semaphore(max_concurrency)
        results = {}

        async def upload_file_async(file_info):
            async with semaphore:
                file_buffer = file_info['file_buffer']
                file_name = file_info['file_name']
                content_type = file_info.get('content_type') or 'application/octet-stream'

                try:
                    # 对于大文件使用分块上传
                    file_size = len(file_buffer)
                    if file_size > 5 * 1024 * 1024:  # 5MB以上使用分块上传
                        url = await self._upload_large_file(
                            s3_client, file_buffer, file_name, content_type
                        )
                    else:
                        # 小文件直接上传
                        s3_client.put_object(
                            Bucket=self.config.bucket_name,
                            Key=file_name,
                            Body=file_buffer,
                            ContentType=content_type
                        )
                        url = self.generate_public_url(file_name)

                    results[file_name] = url
                    print(f"✅ 批量上传成功: {file_name} ({file_size/1024/1024:.2f}MB)")
                except ClientError as e:
                    error_code = e.response.get('Error', {}).get('Code')
                    print(f"❌ 批量上传失败: {file_name}, 错误: {error_code} - {e}")
                    results[file_name] = ''
                except Exception as e:
                    print(f"❌ 批量上传失败: {file_name}, 错误: {e}")
                    results[file_name] = ''

        # 并行上传文件
        tasks = [upload_file_async(file) for file in files]
        await asyncio.gather(*tasks)

        return results

    async def _upload_large_file(
        self,
        s3_client,
        file_buffer: bytes,
        file_name: str,
        content_type: str
    ) -> str:
        """
        分块上传大文件

        参数:
            s3_client: S3客户端
            file_buffer: 文件字节数据
            file_name: 文件名
            content_type: MIME类型

        返回:
            R2 URL
        """
        # 初始化分块上传
        response = s3_client.create_multipart_upload(
            Bucket=self.config.bucket_name,
            Key=file_name,
            ContentType=content_type
        )
        upload_id = response['UploadId']

        # 分块大小：5MB
        part_size = 5 * 1024 * 1024
        file_size = len(file_buffer)
        parts = []

        try:
            # 上传分块
            for i in range(0, file_size, part_size):
                part_number = (i // part_size) + 1
                part_data = file_buffer[i:i+part_size]

                response = s3_client.upload_part(
                    Bucket=self.config.bucket_name,
                    Key=file_name,
                    UploadId=upload_id,
                    PartNumber=part_number,
                    Body=part_data
                )

                parts.append({
                    'PartNumber': part_number,
                    'ETag': response['ETag']
                })

            # 完成分块上传
            s3_client.complete_multipart_upload(
                Bucket=self.config.bucket_name,
                Key=file_name,
                UploadId=upload_id,
                MultipartUpload={'Parts': parts}
            )

            return self.generate_public_url(file_name)
        except Exception as e:
            # 取消分块上传
            s3_client.abort_multipart_upload(
                Bucket=self.config.bucket_name,
                Key=file_name,
                UploadId=upload_id
            )
            raise

    async def batch_upload_thumbnails(
        self,
        thumbnails: List[Dict[str, any]],
        max_concurrency: int = 10  # 增加并发数
    ) -> Dict[str, str]:
        """
        批量上传多个缩略图到 R2

        参数:
            thumbnails: 缩略图列表，每个缩略图包含 {'file_buffer': bytes, 'file_name': str, 'width': Optional[int], 'height': Optional[int]}
            max_concurrency: 最大并发数

        返回:
            字典，键为文件名，值为缩略图R2 URL
        """
        if not self.is_available():
            return {}

        if not thumbnails:
            return {}

        # 获取S3客户端，避免每次上传都创建新连接
        s3_client = self.get_client()
        if not s3_client:
            return {}

        # 限制并发数
        semaphore = asyncio.Semaphore(max_concurrency)
        results = {}

        async def upload_thumbnail_async(thumbnail_info):
            async with semaphore:
                file_buffer = thumbnail_info['file_buffer']
                file_name = thumbnail_info['file_name']
                width = thumbnail_info.get('width', 200)
                height = thumbnail_info.get('height', 200)

                try:
                    # 直接使用S3客户端上传，避免重复创建客户端
                    s3_client.put_object(
                        Bucket=self.config.bucket_name,
                        Key=file_name,
                        Body=file_buffer,
                        ContentType='image/jpeg'
                    )
                    url = self.generate_public_url(file_name)
                    results[file_name] = url
                    print(f"✅ 批量缩略图上传成功: {file_name} ({len(file_buffer)/1024:.2f}KB)")
                except ClientError as e:
                    error_code = e.response.get('Error', {}).get('Code')
                    print(f"❌ 批量缩略图上传失败: {file_name}, 错误: {error_code} - {e}")
                    results[file_name] = ''
                except Exception as e:
                    print(f"❌ 批量缩略图上传失败: {file_name}, 错误: {e}")
                    results[file_name] = ''

        # 并行上传缩略图
        tasks = [upload_thumbnail_async(thumbnail) for thumbnail in thumbnails]
        await asyncio.gather(*tasks)

        return results

    async def upload_thumbnail(
        self,
        file_buffer: bytes,
        file_name: str,
        width: int = 200,
        height: int = 200
    ) -> str:
        """
        上传缩略图到 R2

        参数:
            file_buffer: 文件字节数据
            file_name: 文件名（完整R2键，包含路径）
            width: 缩略图宽度
            height: 缩略图高度

        返回:
            缩略图R2 URL
        """
        if not self.is_available():
            return ''

        s3_client = self.get_client()

        # 生成缩略图
        thumbnail_data, ext = self.generate_thumbnail(file_buffer, width, height)

        # 直接使用传入的文件名作为缩略图R2键
        thumbnail_r2_key = file_name

        # 上传缩略图
        try:
            s3_client.put_object(
                Bucket=self.config.bucket_name,
                Key=thumbnail_r2_key,
                Body=thumbnail_data,
                ContentType='image/jpeg'
            )
            print(f"R2 缩略图上传成功: {thumbnail_r2_key}")
        except ClientError as e:
            print(f"R2 缩略图上传失败: {e}")
            raise

        return self.generate_public_url(thumbnail_r2_key)

    async def delete_file(self, r2_key: str) -> bool:
        """
        从 R2 删除文件

        参数:
            r2_key: R2 对象键

        返回:
            是否删除成功
        """
        print(f"  [R2Delete] 准备删除: {r2_key}")

        if not self.is_available():
            print(f"  [R2Delete] ❌ R2不可用")
            return False

        if not r2_key:
            print(f"  [R2Delete] ❌ R2 Key为空")
            return False

        s3_client = self.get_client()
        if not s3_client:
            print(f"  [R2Delete] ❌ 无法获取R2客户端")
            return False

        try:
            print(f"  [R2Delete] 调用delete_object...")
            response = s3_client.delete_object(
                Bucket=self.config.bucket_name,
                Key=r2_key
            )
            print(f"  [R2Delete] ✅ 删除成功: {r2_key}")
            print(f"  [R2Delete] 响应: {response}")
            return True
        except ClientError as e:
            print(f"  [R2Delete] ❌ 删除失败: {e}")
            print(f"  [R2Delete] 错误代码: {e.response.get('Error', {}).get('Code')}")
            print(f"  [R2Delete] 错误消息: {e.response.get('Error', {}).get('Message')}")
            return False
        except Exception as e:
            print(f"  [R2Delete] ❌ 未知错误: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _get_content_type(self, filename: str) -> str:
        """根据文件名获取 Content-Type"""
        ext = Path(filename).suffix.lower()
        content_types = {
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.png': 'image/png',
            '.gif': 'image/gif',
            '.webp': 'image/webp',
            '.mp4': 'video/mp4',
            '.webm': 'video/webm',
            '.pdf': 'application/pdf'
        }
        return content_types.get(ext, 'application/octet-stream')


# 全局实例
_config_cache = {}
_cache_lock = threading.Lock()


# 向后兼容的方法 - 旧API
class CompatibleR2Uploader:
    """向后兼容的R2上传器包装类"""
    def __init__(self):
        self.uploader = None

    def is_available(self, config: R2Config) -> bool:
        """检查是否可用 - 向后兼容"""
        if not config or not config.enabled:
            return False
        return config.account_id and config.access_key_id and config.secret_access_key

    async def upload_file(self, file_path: str, company_code: str,
                        product_number: str, file_type: str,
                        filename: str, config: R2Config,
                        sub_type: Optional[str] = None,
                        generate_thumbnail: bool = True) -> Dict:
        """旧的上传文件方法 - 向后兼容"""
        import hashlib
        from pathlib import Path

        self.uploader = R2Uploader(config)
        if not self.uploader.is_available():
            return {
                'r2_url': '',
                'thumbnail_url': '',
                'r2_key': '',
                'thumbnail_key': ''
            }

        # 读取文件
        with open(file_path, 'rb') as f:
            file_buffer = f.read()

        # 计算MD5
        hash_md5 = hashlib.md5()
        hash_md5.update(file_buffer)
        file_hash = hash_md5.hexdigest()
        file_extension = Path(filename).suffix.lower()

        # 构建R2键
        r2_key = f"{company_code}/{file_hash}{file_extension}"
        thumbnail_key = f"{company_code}/thumb_{file_hash}{file_extension}"

        # 上传原图
        r2_url = await self.uploader.upload_file(
            file_buffer=file_buffer,
            file_name=r2_key,
            content_type=self.uploader._get_content_type(filename),
            generate_thumbnail=False
        )

        thumbnail_url = ''
        if generate_thumbnail:
            thumbnail_url = await self.uploader.upload_thumbnail(
                file_buffer=file_buffer,
                file_name=thumbnail_key,
                width=200,
                height=200
            )

        return {
            'r2_url': r2_url,
            'thumbnail_url': thumbnail_url,
            'r2_key': r2_key,
            'thumbnail_key': thumbnail_key
        }

    async def delete_file(self, r2_key: str, config: R2Config) -> bool:
        """删除文件 - 向后兼容"""
        if not r2_key:
            return False

        self.uploader = R2Uploader(config)
        if not self.uploader.is_available():
            return False

        return await self.uploader.delete_file(r2_key)

    async def batch_upload_files(
        self,
        files: List[Dict[str, any]],
        config: R2Config,
        max_concurrency: int = 5
    ) -> Dict[str, str]:
        """批量上传多个文件 - 向后兼容"""
        self.uploader = R2Uploader(config)
        if not self.uploader.is_available():
            return {}

        return await self.uploader.batch_upload_files(files, max_concurrency)

    async def batch_upload_thumbnails(
        self,
        thumbnails: List[Dict[str, any]],
        config: R2Config,
        max_concurrency: int = 5
    ) -> Dict[str, str]:
        """批量上传多个缩略图 - 向后兼容"""
        self.uploader = R2Uploader(config)
        if not self.uploader.is_available():
            return {}

        return await self.uploader.batch_upload_thumbnails(thumbnails, max_concurrency)


# 全局向后兼容的实例
_compatible_uploader = None

def get_r2_uploader(config: R2Config = None, company_code: str = "default") -> any:
    """获取 R2 上传器实例 - 支持新旧API"""
    if config is not None:
        # 新API - 返回R2Uploader实例，按公司编码管理
        return R2Uploader(config, company_code)
    else:
        # 旧API - 返回向后兼容的实例
        global _compatible_uploader
        if _compatible_uploader is None:
            _compatible_uploader = CompatibleR2Uploader()
        return _compatible_uploader


def get_cached_config(company_code: str) -> Optional[R2Config]:
    """从缓存获取配置"""
    with _cache_lock:
        config = _config_cache.get(company_code)
        if config and not config.is_expired():
            return config
        return None


def set_cached_config(company_code: str, config: R2Config):
    """设置缓存配置"""
    with _cache_lock:
        _config_cache[company_code] = config


def clear_cached_config(company_code: str):
    """清除缓存配置"""
    with _cache_lock:
        if company_code in _config_cache:
            del _config_cache[company_code]
            print(f"已清除公司 {company_code} 的 R2 配置缓存")