"""
腾讯云COS存储提供者实现

本模块实现了腾讯云COS的具体存储操作，继承自抽象存储接口。
提供完整的COS文件上传、下载、删除等功能。

特性:
- 支持文件、字节数据、PyTorch张量的上传
- 完整的错误处理和重试机制
- 统一的日志记录
- 与S3接口保持一致的API设计
"""

import os
from typing import Optional, Dict, Any, List
import torch
import loguru
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
from ..storage_interface import (
    StorageProvider,
    StorageFactory,
    UploadResult,
    DownloadResult,
)
from ...config_utils import config


class COSStorageProvider(StorageProvider):
    """
    腾讯云COS存储提供者实现

    实现了StorageProvider接口的所有方法，提供完整的COS存储功能。
    采用懒加载模式初始化COS客户端，提高性能。
    """

    def __init__(self, config: Dict[str, Any]):
        """
        初始化COS存储提供者

        Args:
            config: COS配置字典，必须包含secret_id、secret_key和region
        """
        super().__init__(config)
        self.bucket_name = config.get("bucket_name", "bwkj-cos-1324682537")
        self.region = config["region"]
        self.secret_id = config["secret_id"]
        self.secret_key = config["secret_key"]
        self._client = None

    def _validate_config(self) -> None:
        """
        验证COS配置的完整性

        Raises:
            ValueError: 配置信息缺失时抛出异常
        """
        required_keys = ["secret_id", "secret_key", "region"]
        missing_keys = [key for key in required_keys if not self.config.get(key)]

        if missing_keys:
            raise ValueError(
                f"COS配置缺失必要参数: {missing_keys}. " f"请检查配置文件或环境变量"
            )

    @property
    def client(self):
        """
        获取COS客户端实例（懒加载模式）

        Returns:
            CosS3Client: COS客户端实例
        """
        if self._client is None:
            try:
                cos_config = CosConfig(
                    Region=self.region,
                    SecretId=self.secret_id,
                    SecretKey=self.secret_key,
                )
                self._client = CosS3Client(cos_config)
                loguru.logger.info(f"COS客户端初始化成功，区域: {self.region}")
            except Exception as e:
                loguru.logger.error(f"COS客户端初始化失败: {e}")
                raise

        return self._client

    def upload_file(
        self,
        local_path: str,
        remote_key: str,
        content_type: Optional[str] = None,
        **kwargs,
    ) -> UploadResult:
        """
        上传本地文件到COS

        Args:
            local_path: 本地文件路径
            remote_key: COS中的键名
            content_type: 文件内容类型
            **kwargs: 额外的上传参数

        Returns:
            UploadResult: 上传操作结果
        """
        try:
            if not os.path.exists(local_path):
                error_msg = f"本地文件不存在: {local_path}"
                loguru.logger.error(error_msg)
                return UploadResult(
                    success=False,
                    key=remote_key,
                    message=error_msg,
                    error=FileNotFoundError(error_msg),
                )

            file_size = os.path.getsize(local_path)
            loguru.logger.info(
                f"开始上传文件到COS: {local_path} -> cos://{self.bucket_name}/{remote_key} "
                f"({file_size} bytes)"
            )

            # 执行上传操作，包含重试机制
            for attempt in range(3):  # 最多重试3次
                try:
                    response = self.client.upload_file(
                        Bucket=self.bucket_name,
                        Key=remote_key,
                        LocalFilePath=local_path,
                        **kwargs,
                    )
                    break
                except (CosClientError, CosServiceError) as e:
                    if attempt == 2:  # 最后一次尝试失败
                        raise e
                    loguru.logger.warning(
                        f"COS上传尝试 {attempt + 1} 失败，重试中: {e}"
                    )

            success_msg = f"文件上传成功: cos://{self.bucket_name}/{remote_key}"
            loguru.logger.info(success_msg)

            return UploadResult(
                success=True, key=remote_key, size=file_size, message=success_msg
            )

        except Exception as e:
            error_msg = f"COS文件上传失败: {str(e)}"
            loguru.logger.error(f"{error_msg} (文件: {local_path}, 键: {remote_key})")
            return UploadResult(
                success=False, key=remote_key, message=error_msg, error=e
            )

    def upload_bytes(
        self, data: bytes, remote_key: str, content_type: Optional[str] = None, **kwargs
    ) -> UploadResult:
        """
        上传字节数据到COS

        Args:
            data: 字节数据
            remote_key: COS中的键名
            content_type: 文件内容类型
            **kwargs: 额外的上传参数

        Returns:
            UploadResult: 上传操作结果
        """
        try:
            loguru.logger.info(
                f"开始上传字节数据到COS: cos://{self.bucket_name}/{remote_key} "
                f"({len(data)} bytes)"
            )

            # 执行上传操作，包含重试机制
            for attempt in range(3):
                try:
                    response = self.client.put_object(
                        Bucket=self.bucket_name,
                        Key=remote_key,
                        Body=data,
                        ContentType=content_type,
                        **kwargs,
                    )
                    break
                except (CosClientError, CosServiceError) as e:
                    if attempt == 2:
                        raise e
                    loguru.logger.warning(
                        f"COS字节数据上传尝试 {attempt + 1} 失败，重试中: {e}"
                    )

            success_msg = f"字节数据上传成功: cos://{self.bucket_name}/{remote_key}"
            loguru.logger.info(success_msg)

            return UploadResult(
                success=True, key=remote_key, size=len(data), message=success_msg
            )

        except Exception as e:
            error_msg = f"COS字节数据上传失败: {str(e)}"
            loguru.logger.error(
                f"{error_msg} (键: {remote_key}, 大小: {len(data)} bytes)"
            )
            return UploadResult(
                success=False, key=remote_key, message=error_msg, error=e
            )

    def upload_tensor(
        self, tensor: torch.Tensor, remote_key: str, format: str = "PNG", **kwargs
    ) -> UploadResult:
        """
        上传PyTorch张量作为图像到COS

        Args:
            tensor: PyTorch张量
            remote_key: COS中的键名
            format: 图像格式（PNG, JPEG等）
            **kwargs: 额外的上传参数

        Returns:
            UploadResult: 上传操作结果
        """
        try:
            from ...image_utils import tensor_to_tempfile

            loguru.logger.info(
                f"开始上传张量到COS: cos://{self.bucket_name}/{remote_key} "
                f"(形状: {tensor.shape}, 格式: {format})"
            )

            # 将张量转换为临时文件
            temp_file = tensor_to_tempfile(tensor, format=format)
            temp_path = temp_file.name

            try:
                # 设置内容类型
                content_type = f"image/{format.lower()}"
                if format.upper() == "JPEG":
                    content_type = "image/jpeg"

                # 上传临时文件
                result = self.upload_file(
                    temp_path, remote_key, content_type=content_type, **kwargs
                )

                if result.success:
                    loguru.logger.info(
                        f"张量上传成功: cos://{self.bucket_name}/{remote_key}"
                    )

                return result

            finally:
                # 清理临时文件
                if os.path.exists(temp_path):
                    os.unlink(temp_path)
                    loguru.logger.debug(f"临时文件已清理: {temp_path}")

        except Exception as e:
            error_msg = f"COS张量上传失败: {str(e)}"
            loguru.logger.error(
                f"{error_msg} (键: {remote_key}, 张量形状: {tensor.shape})"
            )
            return UploadResult(
                success=False, key=remote_key, message=error_msg, error=e
            )

    def download_file(
        self, remote_key: str, local_path: str, **kwargs
    ) -> DownloadResult:
        """
        从COS下载文件到本地

        Args:
            remote_key: COS中的键名
            local_path: 本地保存路径
            **kwargs: 额外的下载参数

        Returns:
            DownloadResult: 下载操作结果
        """
        try:
            loguru.logger.info(
                f"开始从COS下载文件: cos://{self.bucket_name}/{remote_key} -> {local_path}"
            )

            # 确保本地目录存在
            os.makedirs(os.path.dirname(local_path), exist_ok=True)

            # 检查文件是否存在
            if not self.file_exists(remote_key):
                error_msg = f"COS文件不存在: cos://{self.bucket_name}/{remote_key}"
                loguru.logger.error(error_msg)
                return DownloadResult(
                    success=False, message=error_msg, error=FileNotFoundError(error_msg)
                )

            # 执行下载操作，包含重试机制
            for attempt in range(3):
                try:
                    response = self.client.download_file(
                        Bucket=self.bucket_name,
                        Key=remote_key,
                        DestFilePath=local_path,
                        **kwargs,
                    )
                    break
                except (CosClientError, CosServiceError) as e:
                    if attempt == 2:
                        raise e
                    loguru.logger.warning(
                        f"COS下载尝试 {attempt + 1} 失败，重试中: {e}"
                    )

            success_msg = f"文件下载成功: {local_path}"
            loguru.logger.info(success_msg)

            return DownloadResult(
                success=True, local_path=local_path, message=success_msg
            )

        except Exception as e:
            error_msg = f"COS文件下载失败: {str(e)}"
            loguru.logger.error(
                f"{error_msg} (键: {remote_key}, 本地路径: {local_path})"
            )
            return DownloadResult(success=False, message=error_msg, error=e)

    def download_bytes(self, remote_key: str, **kwargs) -> DownloadResult:
        """
        从COS下载文件为字节数据

        Args:
            remote_key: COS中的键名
            **kwargs: 额外的下载参数

        Returns:
            DownloadResult: 下载操作结果，数据包含在data字段中
        """
        try:
            loguru.logger.info(
                f"开始从COS下载字节数据: cos://{self.bucket_name}/{remote_key}"
            )

            # 检查文件是否存在
            if not self.file_exists(remote_key):
                error_msg = f"COS文件不存在: cos://{self.bucket_name}/{remote_key}"
                loguru.logger.error(error_msg)
                return DownloadResult(
                    success=False, message=error_msg, error=FileNotFoundError(error_msg)
                )

            # 执行下载操作，包含重试机制
            for attempt in range(3):
                try:
                    response = self.client.get_object(
                        Bucket=self.bucket_name, Key=remote_key, **kwargs
                    )
                    data = response["Body"].read()
                    break
                except (CosClientError, CosServiceError) as e:
                    if attempt == 2:
                        raise e
                    loguru.logger.warning(
                        f"COS字节数据下载尝试 {attempt + 1} 失败，重试中: {e}"
                    )

            success_msg = f"字节数据下载成功: {len(data)} bytes"
            loguru.logger.info(success_msg)

            return DownloadResult(success=True, data=data, message=success_msg)

        except Exception as e:
            error_msg = f"COS字节数据下载失败: {str(e)}"
            loguru.logger.error(f"{error_msg} (键: {remote_key})")
            return DownloadResult(success=False, message=error_msg, error=e)

    def delete_file(self, remote_key: str, **kwargs) -> bool:
        """
        删除COS中的文件

        Args:
            remote_key: COS中的键名
            **kwargs: 额外的删除参数

        Returns:
            bool: 删除是否成功
        """
        try:
            loguru.logger.info(f"删除COS文件: cos://{self.bucket_name}/{remote_key}")

            # 执行删除操作，包含重试机制
            for attempt in range(3):
                try:
                    self.client.delete_object(
                        Bucket=self.bucket_name, Key=remote_key, **kwargs
                    )
                    break
                except (CosClientError, CosServiceError) as e:
                    if attempt == 2:
                        raise e
                    loguru.logger.warning(
                        f"COS删除尝试 {attempt + 1} 失败，重试中: {e}"
                    )

            loguru.logger.info(f"文件删除成功: cos://{self.bucket_name}/{remote_key}")
            return True

        except Exception as e:
            loguru.logger.error(f"COS文件删除失败: {e} (键: {remote_key})")
            return False

    def file_exists(self, remote_key: str, **kwargs) -> bool:
        """
        检查COS中文件是否存在

        Args:
            remote_key: COS中的键名
            **kwargs: 额外的检查参数

        Returns:
            bool: 文件是否存在
        """
        try:
            self.client.head_object(Bucket=self.bucket_name, Key=remote_key, **kwargs)
            return True
        except Exception:
            return False

    def list_files(
        self, prefix: str = "", max_keys: int = 1000, **kwargs
    ) -> List[Dict[str, Any]]:
        """
        列出COS中的文件

        Args:
            prefix: 文件前缀过滤
            max_keys: 最大返回数量
            **kwargs: 额外的列表参数

        Returns:
            List[Dict[str, Any]]: 文件信息列表
        """
        try:
            loguru.logger.info(
                f"列出COS文件: cos://{self.bucket_name}/{prefix} (最大: {max_keys})"
            )

            response = self.client.list_objects(
                Bucket=self.bucket_name, Prefix=prefix, MaxKeys=max_keys, **kwargs
            )

            files = []
            if "Contents" in response:
                for obj in response["Contents"]:
                    files.append(
                        {
                            "key": obj["Key"],
                            "size": obj["Size"],
                            "last_modified": obj["LastModified"],
                            "etag": obj.get("ETag", "").strip('"'),
                            "storage_class": obj.get("StorageClass", "STANDARD"),
                        }
                    )

            loguru.logger.info(f"找到 {len(files)} 个文件")
            return files

        except Exception as e:
            loguru.logger.error(f"COS文件列表获取失败: {e} (前缀: {prefix})")
            return []

    def get_file_url(self, remote_key: str, expires_in: int = 3600, **kwargs) -> str:
        """
        获取COS文件的访问URL

        Args:
            remote_key: COS中的键名
            expires_in: URL过期时间（秒）
            **kwargs: 额外的URL生成参数

        Returns:
            str: 文件访问URL
        """
        try:
            # 生成预签名URL
            url = self.client.get_presigned_url(
                Method="GET",
                Bucket=self.bucket_name,
                Key=remote_key,
                Expired=expires_in,
                **kwargs,
            )

            return url

        except Exception as e:
            loguru.logger.error(f"COS URL生成失败: {e} (键: {remote_key})")
            raise


class COSStorageFactory(StorageFactory):
    """
    COS存储工厂实现

    负责创建COS存储提供者实例，处理配置验证和初始化。
    """

    def create_provider(self, config_dict: Dict[str, Any]) -> StorageProvider:
        """
        创建COS存储提供者实例

        Args:
            config_dict: COS配置字典

        Returns:
            StorageProvider: COS存储提供者实例
        """
        # 如果没有提供配置，尝试从全局配置获取
        if not config_dict or not config_dict.get("secret_id"):
            if config.has_cos_config():
                cos_config = config.get_cos_config()
                config_dict = {
                    "secret_id": cos_config["secret_id"],
                    "secret_key": cos_config["secret_key"],
                    "region": cos_config["region"],
                    "bucket_name": cos_config.get("bucket_name"),
                    **config_dict,
                }
            else:
                raise ValueError("未提供有效的COS配置，且全局配置中也没有COS配置")

        return COSStorageProvider(config_dict)

    def get_supported_types(self) -> List[str]:
        """
        获取支持的存储类型列表

        Returns:
            List[str]: 支持的存储类型
        """
        return ["cos", "qcloud", "tencent"]