"""
S3 兼容存储模块
提供对象存储的上传、下载、删除等功能
"""

import os
import re
import logging
from pathlib import Path
from typing import Optional, Any, Dict, List, Iterable
from uuid import uuid4

import boto3
from botocore.exceptions import ClientError
from boto3.s3.transfer import TransferConfig

from .models import ListFilesResult

logger = logging.getLogger(__name__)

# 允许的文件名字符集（面向用户输入的约束）
FILE_NAME_ALLOWED_RE = re.compile(r"^[A-Za-z0-9._\-/]+$")


def _load_env() -> None:
    """加载环境变量"""
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass

    try:
        from coze_workload_identity import Client
        client = Client()
        env_vars = client.get_project_env_vars()
        client.close()
        for env_var in env_vars:
            if env_var.key not in os.environ:
                os.environ[env_var.key] = env_var.value
    except Exception as e:
        logger.debug(f"coze_workload_identity not available: {e}")


class S3SyncStorage:
    """S3 兼容存储实现"""

    def __init__(
        self,
        *,
        endpoint_url: Optional[str] = None,
        access_key: str = "",
        secret_key: str = "",
        bucket_name: Optional[str] = None,
        region: str = "cn-beijing",
    ):
        """
        初始化 S3 存储客户端

        Args:
            endpoint_url: S3 兼容端点，默认从环境变量 COZE_BUCKET_ENDPOINT_URL 获取
            access_key: 访问密钥，默认为空
            secret_key: 秘密密钥，默认为空
            bucket_name: 桶名称，默认从环境变量 COZE_BUCKET_NAME 获取
            region: 区域，默认 cn-beijing
        """
        _load_env()
        self.endpoint_url = os.environ.get("COZE_BUCKET_ENDPOINT_URL") or endpoint_url or ""
        self.access_key = access_key
        self.secret_key = secret_key
        self.bucket_name = bucket_name or os.environ.get("COZE_BUCKET_NAME") or ""
        self.region = region
        self._client = None

    def _get_client(self):
        if self._client is None:
            endpoint = self.endpoint_url
            if not endpoint:
                try:
                    from coze_workload_identity import Client as CozeEnvClient
                    coze_env_client = CozeEnvClient()
                    env_vars = coze_env_client.get_project_env_vars()
                    coze_env_client.close()
                    for env_var in env_vars:
                        if env_var.key == "COZE_BUCKET_ENDPOINT_URL":
                            endpoint = env_var.value.replace("'", "'\\''")
                            self.endpoint_url = endpoint
                            break
                except Exception as e:
                    logger.error(f"Error loading COZE_BUCKET_ENDPOINT_URL: {e}")

            if not endpoint:
                logger.error("未配置存储端点：请设置 endpoint_url")
                raise ValueError("未配置存储端点：请设置 endpoint_url")

            client = boto3.client(
                "s3",
                endpoint_url=endpoint,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            )

            # 注册 before-call 钩子，发送前注入 x-storage-token 头
            def _inject_header(**kwargs):
                try:
                    from coze_workload_identity import Client as CozeClient
                    coze_client = CozeClient()
                    try:
                        token = coze_client.get_access_token()
                    except Exception as e:
                        logger.error("Error loading COZE_WORKLOAD_IDENTITY_TOKEN: %s", e)
                        token = None
                        raise e
                    finally:
                        coze_client.close()
                    params = kwargs.get("params", {})
                    headers = params.setdefault("headers", {})
                    headers["x-storage-token"] = token
                except Exception as e:
                    logger.error("Error loading COZE_WORKLOAD_IDENTITY_TOKEN: %s", e)

            client.meta.events.register("before-call.s3", _inject_header)
            self._client = client
        return self._client

    def _generate_object_key(self, *, original_name: str) -> str:
        suffix = Path(original_name).suffix.lower()
        stem = Path(original_name).stem
        uniq = uuid4().hex[:8]
        return f"{stem}_{uniq}{suffix}"

    def _extract_logid(self, e: Exception) -> Optional[str]:
        """从 ClientError 中提取 x-tt-logid"""
        if isinstance(e, ClientError):
            headers = (e.response or {}).get("ResponseMetadata", {}).get("HTTPHeaders", {})
            return headers.get("x-tt-logid")
        return None

    def _error_msg(self, msg: str, e: Exception) -> str:
        """构建带 logid 的错误信息"""
        logid = self._extract_logid(e)
        if logid:
            return f"{msg}: {e} (x-tt-logid: {logid})"
        return f"{msg}: {e}"

    def _resolve_bucket(self, bucket: Optional[str]) -> str:
        """统一解析 bucket 来源，确保得到有效桶名"""
        target_bucket = bucket or os.environ.get("COZE_BUCKET_NAME") or self.bucket_name
        if not target_bucket:
            raise ValueError("未配置 bucket：请传入 bucket 或设置 COZE_BUCKET_NAME，或在实例化时提供 bucket_name")
        return target_bucket

    def _validate_file_name(self, name: str) -> None:
        """校验 S3 对象命名：长度≤1024；允许 [A-Za-z0-9._-/]；不以 / 起止且不含 //"""
        msg = (
            "file name invalid: 文件名需满足以下 S3 对象命名规范："
            "1) 长度 1–1024 字节；"
            "2) 仅允许字母、数字、点(.)、下划线(_)、短横(-)、目录分隔符(/)；"
            "3) 不允许空格或以下特殊字符：? # & % { } ^ [ ] ` \\ < > ~ | \" ' + = : ;；"
            "4) 不以 / 开头或结尾，且不包含连续的 //；"
            "示例：report_2025-12-11.pdf、images/photo-01.png。"
        )

        if not name or not name.strip():
            raise ValueError(msg + "（原因：文件名为空）")

        if len(name.encode("utf-8")) > 1024:
            raise ValueError(msg + "（原因：长度超过 1024 字节）")

        if name.startswith("/") or name.endswith("/"):
            raise ValueError(msg + "（原因：以 / 开头或结尾）")
        if "//" in name:
            raise ValueError(msg + "（原因：包含连续的 //）")

        if not FILE_NAME_ALLOWED_RE.match(name):
            bad = re.findall(r"[^A-Za-z0-9._\-/]", name)
            example = bad[0] if bad else "非法字符"
            raise ValueError(msg + f"（原因：包含非法字符，例如：{example}）")

    def upload_file(
        self,
        *,
        file_content: bytes,
        file_name: str,
        content_type: str = "application/octet-stream",
        bucket: Optional[str] = None,
    ) -> str:
        """
        上传文件

        Args:
            file_content: 文件内容
            file_name: 原始文件名
            content_type: MIME 类型
            bucket: 目标桶（可选）

        Returns:
            str: 对象 key
        """
        self._validate_file_name(file_name)
        try:
            client = self._get_client()
            object_key = self._generate_object_key(original_name=file_name)
            target_bucket = self._resolve_bucket(bucket)
            client.put_object(Bucket=target_bucket, Key=object_key, Body=file_content, ContentType=content_type)
            return object_key
        except Exception as e:
            logger.error(self._error_msg("Error uploading file to S3", e))
            raise e

    def delete_file(self, *, file_key: str, bucket: Optional[str] = None) -> bool:
        """
        删除文件

        Args:
            file_key: 对象 key
            bucket: 目标桶（可选）

        Returns:
            bool: 是否成功
        """
        try:
            client = self._get_client()
            target_bucket = self._resolve_bucket(bucket)
            client.delete_object(Bucket=target_bucket, Key=file_key)
            return True
        except Exception as e:
            logger.error(self._error_msg("Error deleting file from S3", e))
            raise e

    def file_exists(self, *, file_key: str, bucket: Optional[str] = None) -> bool:
        """
        检查文件是否存在

        Args:
            file_key: 对象 key
            bucket: 目标桶（可选）

        Returns:
            bool: 是否存在
        """
        try:
            client = self._get_client()
            target_bucket = self._resolve_bucket(bucket)
            client.head_object(Bucket=target_bucket, Key=file_key)
            return True
        except ClientError as e:
            code = (e.response or {}).get("Error", {}).get("Code", "")
            if code in {"404", "NoSuchKey", "NotFound"}:
                return False
            logger.error(self._error_msg("Error checking file existence in S3", e))
            return False
        except Exception as e:
            logger.error(self._error_msg("Error checking file existence in S3", e))
            return False

    def read_file(self, *, file_key: str, bucket: Optional[str] = None) -> bytes:
        """
        读取文件

        Args:
            file_key: 对象 key
            bucket: 目标桶（可选）

        Returns:
            bytes: 文件内容
        """
        try:
            client = self._get_client()
            target_bucket = self._resolve_bucket(bucket)
            resp = client.get_object(Bucket=target_bucket, Key=file_key)
            body = resp.get("Body")
            if body is None:
                raise RuntimeError("S3 get_object returned no Body")
            try:
                return body.read()
            finally:
                try:
                    body.close()
                except Exception as ce:
                    logger.debug("Failed to close S3 response body: %s", ce)
        except Exception as e:
            logger.error(self._error_msg("Error reading file from S3", e))
            raise e

    def list_files(
        self,
        *,
        prefix: Optional[str] = None,
        bucket: Optional[str] = None,
        max_keys: int = 1000,
        continuation_token: Optional[str] = None,
    ) -> ListFilesResult:
        """
        列出对象

        Args:
            prefix: 前缀过滤
            bucket: 目标桶（可选）
            max_keys: 最大返回数量（1-1000）
            continuation_token: 分页 token

        Returns:
            ListFilesResult: 包含 keys, is_truncated, next_continuation_token
        """
        try:
            client = self._get_client()
            target_bucket = self._resolve_bucket(bucket)
            if max_keys <= 0 or max_keys > 1000:
                raise ValueError("max_keys 必须在 1 到 1000 之间")

            kwargs: Dict[str, Any] = {
                "Bucket": target_bucket,
                "MaxKeys": max_keys,
                "Prefix": prefix,
                "ContinuationToken": continuation_token,
            }
            kwargs = {k: v for k, v in kwargs.items() if v is not None}

            resp = client.list_objects_v2(**kwargs)
            contents = resp.get("Contents", []) or []
            keys: List[str] = [item.get("Key") for item in contents if isinstance(item, dict) and item.get("Key")]
            return {
                "keys": keys,
                "is_truncated": bool(resp.get("IsTruncated")),
                "next_continuation_token": resp.get("NextContinuationToken"),
            }
        except ClientError as e:
            code = (e.response or {}).get("Error", {}).get("Code", "")
            logger.error(self._error_msg(f"Error listing files in S3 (code={code})", e))
            raise e
        except Exception as e:
            logger.error(self._error_msg("Error listing files in S3", e))
            raise e

    def generate_presigned_url(
        self,
        *,
        key: str,
        bucket: Optional[str] = None,
        expire_time: int = 86400,
    ) -> str:
        """
        生成签名 URL

        Args:
            key: 对象 key
            bucket: 目标桶（可选）
            expire_time: 过期时间（秒），默认 86400

        Returns:
            str: 签名 URL
        """
        import json
        import urllib.request as urllib_request

        try:
            from coze_workload_identity import Client as CozeClient
            coze_client = CozeClient()
            try:
                token = coze_client.get_access_token()
            finally:
                try:
                    coze_client.close()
                except Exception:
                    pass
        except Exception as e:
            logger.error(f"Error loading x-storage-token: {e}")
            raise RuntimeError(f"获取 x-storage-token 失败: {e}")

        try:
            sign_base = os.environ.get("COZE_BUCKET_ENDPOINT_URL") or self.endpoint_url
            if not sign_base:
                raise ValueError("未配置签名端点：请设置 COZE_BUCKET_ENDPOINT_URL 或传入 endpoint_url")
            sign_url_endpoint = sign_base.rstrip("/") + "/sign-url"

            headers = {
                "Content-Type": "application/json",
                "x-storage-token": token,
            }

            target_bucket = self._resolve_bucket(bucket)
            payload = {"bucket_name": target_bucket, "path": key, "expire_time": expire_time}
            data = json.dumps(payload).encode("utf-8")
            request = urllib_request.Request(sign_url_endpoint, data=data, headers=headers, method="POST")
        except Exception as e:
            logger.error(f"Error creating request for sign-url: {e}")
            raise RuntimeError(f"创建 sign-url 请求失败: {e}")

        try:
            with urllib_request.urlopen(request) as resp:
                resp_bytes = resp.read()
                content_type = resp.headers.get("Content-Type", "")
                text = resp_bytes.decode("utf-8", errors="replace")
                if "application/json" in content_type or text.strip().startswith("{"):
                    try:
                        obj = json.loads(text)
                    except Exception:
                        return text
                    data = obj.get("data")
                    if isinstance(data, dict) and "url" in data:
                        return data["url"]
                    url_value = obj.get("url") or obj.get("signed_url") or obj.get("presigned_url")
                    if url_value:
                        return url_value
                    raise ValueError("签名服务返回缺少 data.url/url 字段")
                return text
        except Exception as e:
            raise RuntimeError(f"生成签名URL失败: {e}")

    def stream_upload_file(
        self,
        *,
        fileobj,
        file_name: str,
        content_type: str = "application/octet-stream",
        bucket: Optional[str] = None,
        multipart_chunksize: int = 5 * 1024 * 1024,
        multipart_threshold: int = 5 * 1024 * 1024,
        max_concurrency: int = 1,
        use_threads: bool = False,
    ) -> str:
        """
        流式上传文件对象

        Args:
            fileobj: 文件对象（需有 read() 方法）
            file_name: 原始文件名
            content_type: MIME 类型
            bucket: 目标桶（可选）
            multipart_chunksize: 分片大小（默认 5MB）
            multipart_threshold: 触发分片上传的阈值（默认 5MB）
            max_concurrency: 并发数（默认 1）
            use_threads: 是否启用多线程（默认 False）

        Returns:
            str: 对象 key
        """
        try:
            client = self._get_client()
            target_bucket = self._resolve_bucket(bucket)
            key = self._generate_object_key(original_name=file_name)

            extra_args = {"ContentType": content_type} if content_type else {}
            config = TransferConfig(
                multipart_chunksize=multipart_chunksize,
                multipart_threshold=multipart_threshold,
                max_concurrency=max_concurrency,
                use_threads=use_threads,
            )
            client.upload_fileobj(Fileobj=fileobj, Bucket=target_bucket, Key=key, ExtraArgs=extra_args, Config=config)
            return key
        except Exception as e:
            logger.error(self._error_msg("Error streaming upload (fileobj) to S3", e))
            raise e

    def upload_from_url(
        self,
        *,
        url: str,
        bucket: Optional[str] = None,
        timeout: int = 30,
    ) -> str:
        """
        从 URL 下载并上传到 S3

        Args:
            url: 源文件 URL
            bucket: 目标桶（可选）
            timeout: HTTP 超时（秒），默认 30

        Returns:
            str: 对象 key
        """
        import urllib.request as urllib_request
        from urllib.parse import urlparse, unquote

        try:
            request = urllib_request.Request(url)
            with urllib_request.urlopen(request, timeout=timeout) as resp:
                parsed = urlparse(url)
                file_name = Path(unquote(parsed.path)).name or "file"
                content_type = resp.headers.get("Content-Type", "application/octet-stream")
                return self.stream_upload_file(
                    fileobj=resp,
                    file_name=file_name,
                    content_type=content_type,
                    bucket=bucket,
                )
        except Exception as e:
            logger.error(self._error_msg("Error uploading from URL to S3", e))
            raise e

    def trunk_upload_file(
        self,
        *,
        chunk_iter: Iterable[bytes],
        file_name: str,
        content_type: str = "application/octet-stream",
        bucket: Optional[str] = None,
        part_size: int = 5 * 1024 * 1024,
    ) -> str:
        """
        分块流式上传

        Args:
            chunk_iter: 字节块迭代器
            file_name: 原始文件名
            content_type: MIME 类型
            bucket: 目标桶（可选）
            part_size: 分片大小（默认 5MB）

        Returns:
            str: 对象 key
        """
        client = self._get_client()
        target_bucket = self._resolve_bucket(bucket)
        key = self._generate_object_key(original_name=file_name)

        # 初始化分片上传
        try:
            init_resp = client.create_multipart_upload(Bucket=target_bucket, Key=key, ContentType=content_type)
            upload_id = init_resp["UploadId"]
        except Exception as e:
            logger.error(self._error_msg("create_multipart_upload failed", e))
            raise e

        parts = []
        part_number = 1
        buffer = bytearray()
        try:
            for chunk in chunk_iter:
                if not chunk:
                    continue
                buffer.extend(chunk)
                while len(buffer) >= part_size:
                    data = bytes(buffer[:part_size])
                    buffer = buffer[part_size:]
                    resp = client.upload_part(
                        Bucket=target_bucket,
                        Key=key,
                        UploadId=upload_id,
                        PartNumber=part_number,
                        Body=data,
                    )
                    parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
                    part_number += 1

            # 上传最后不足 part_size 的余量
            if len(buffer) > 0:
                resp = client.upload_part(
                    Bucket=target_bucket,
                    Key=key,
                    UploadId=upload_id,
                    PartNumber=part_number,
                    Body=bytes(buffer),
                )
                parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})

            # 完成分片
            client.complete_multipart_upload(
                Bucket=target_bucket,
                Key=key,
                UploadId=upload_id,
                MultipartUpload={"Parts": parts},
            )
            return key
        except Exception as e:
            logger.error(self._error_msg("multipart upload failed", e))
            try:
                client.abort_multipart_upload(Bucket=target_bucket, Key=key, UploadId=upload_id)
            except Exception as ae:
                logger.error(self._error_msg("abort_multipart_upload failed", ae))
            raise e


__all__ = [
    "S3SyncStorage",
]
