from pathlib import Path
from typing import Optional

from qcloud_cos import (
    CosConfig,
    CosS3Client,
    CosServiceError,
)


class LightCOSClient:
    """
    腾讯云轻量COS客户端
    """

    def __init__(
        self, region: str, secret_id: str, secret_key: str, bucket: str
    ) -> None:
        self.client = CosS3Client(
            conf=CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
        )
        self.bucket = bucket
        self.total_bytes_transferred = 0

    def upload_file(self, key: str, content: bytes) -> None:
        """
        上传文件
        @param key: 文件key
        @param content: 文件内容
        """
        self.client.put_object(
            Bucket=self.bucket,
            Body=content,
            Key=key,
        )

    def get_file_etag(self, key: str) -> Optional[str]:
        """
        获取文件ETag
        @param key: 文件key
        @return: 文件ETag
        """
        try:
            response = self.client.head_object(Bucket=self.bucket, Key=key)
            # HEAD请求不传输文件内容，只传输元数据，字节数很小，可以忽略
            return response.get("ETag", response.get("Etag", "")).strip('"')
        except CosServiceError as e:
            if e.get_error_code() == "NoSuchResource":
                return None
            raise

    def download_file(self, key: str, local_path: Path) -> int:
        """
        下载文件（原子写：先写临时文件再 rename，避免下载过程中计算进程读到不完整数据）
        @param key: 文件key
        @param local_path: 本地保存路径
        @return: 下载的字节数
        """
        import os
        # 确保文件路径存在
        local_path.parent.mkdir(parents=True, exist_ok=True)
        # 获取文件
        response = self.client.get_object(Bucket=self.bucket, Key=key)
        # 获取文件大小
        content_length = response.get("Content-Length", "0")
        try:
            bytes_transferred = int(content_length)
        except ValueError:
            bytes_transferred = 0
        # 原子写：先写临时文件，再 rename 替换
        tmp_path = local_path.with_suffix(".cos_tmp")
        with open(tmp_path, "wb") as f:
            data = response["Body"].get_raw_stream().read()
            f.write(data)
            if bytes_transferred == 0:
                bytes_transferred = len(data)
        os.rename(tmp_path, local_path)
        # 更新总字节数
        self.total_bytes_transferred += bytes_transferred
        return bytes_transferred

    def exists_file(self, key: str) -> bool:
        """
        检查文件是否存在
        @param key: 文件key
        @return: 文件是否存在
        """
        return self.get_file_etag(key) is not None

    def list_objects(self, prefix: str = "") -> list[str]:
        """
        列出指定前缀的对象
        @param prefix: 前缀
        @return: 对象名称列表
        """
        objects = []
        marker = ""
        while True:
            response = self.client.list_objects(
                Bucket=self.bucket, Prefix=prefix, Marker=marker, MaxKeys=1000
            )
            if "Contents" in response:
                objects.extend([obj["Key"] for obj in response["Contents"]])
            if response.get("IsTruncated") == "false":
                break
            marker = response.get("NextMarker", "")
        return objects

    def reset_bytes_transferred(self) -> None:
        """
        重置总字节数
        """
        self.total_bytes_transferred = 0

    def get_bytes_transferred(self) -> int:
        """
        获取总字节数
        @return: 总字节数
        """
        return self.total_bytes_transferred
