#!/usr/bin/env python3
"""
京东云对象存储存量迁移脚本。

遍历本地 upload.base-dir（默认 /data/cfc-uploads）下的所有文件，
上传到京东云 OSS bucket cfc，保持原有目录结构。

用法:
    # 查看会迁移哪些文件（不实际上传）
    python3 scripts/migrate-to-jdcloud.py --dry-run

    # 执行迁移（自动跳过云端已存在的文件）
    python3 scripts/migrate-to-jdcloud.py

    # 强制覆盖云端已存在的文件
    python3 scripts/migrate-to-jdcloud.py --overwrite

    # 指定本地目录
    python3 scripts/migrate-to-jdcloud.py --local-dir /data/cfc-uploads
"""

import os
import sys
import boto3
from botocore.config import Config
from pathlib import Path
import argparse
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("migrate")

# 京东云 OSS 配置
ENDPOINT = "s3.cn-north-1.jdcloud-oss.com"
ACCESS_KEY = "JDC_F351EC3CB1F8593204ABCDA2CBF2"
SECRET_KEY = "970C0C6985A6BF29D0C79AE3693B9899"
BUCKET = "cfc"
REGION = "cn-north-1"

# 公网URL前缀
PUBLIC_URL_PREFIX = f"https://{BUCKET}.{ENDPOINT}/"


def get_client():
    return boto3.client(
        "s3",
        endpoint_url=f"https://{ENDPOINT}",
        aws_access_key_id=ACCESS_KEY,
        aws_secret_access_key=SECRET_KEY,
        config=Config(s3={"addressing_style": "path"}, signature_version="s3v4"),
        region_name=REGION,
    )


def list_cloud_keys(client):
    """列出云端已存在的所有 key"""
    keys = set()
    marker = None
    while True:
        if marker:
            resp = client.list_objects_v2(Bucket=BUCKET, MaxKeys=1000, StartAfter=marker)
        else:
            resp = client.list_objects_v2(Bucket=BUCKET, MaxKeys=1000)
        if "Contents" not in resp:
            break
        for obj in resp["Contents"]:
            keys.add(obj["Key"])
            marker = obj["Key"]
        if not resp.get("IsTruncated"):
            break
    return keys


def upload_file(client, local_path, key, dry_run=False, overwrite=False, cloud_keys=None):
    """上传单个文件"""
    target_url = PUBLIC_URL_PREFIX + key

    if cloud_keys and key in cloud_keys and not overwrite:
        log.info("  [SKIP] 已在云端存在: %s", target_url)
        return False

    if dry_run:
        log.info("  [DRY-RUN] 将上传: %s -> %s", local_path, target_url)
        return False

    try:
        content_type = guess_content_type(local_path)
        extra_args = {"ContentType": content_type} if content_type else {}
        client.upload_file(str(local_path), BUCKET, key, ExtraArgs=extra_args)
        log.info("  [OK] %s", target_url)
        return True
    except Exception as e:
        log.error("  [FAIL] %s: %s", local_path, e)
        return False


def guess_content_type(path):
    """根据扩展名猜测 Content-Type"""
    ext = Path(path).suffix.lower()
    mapping = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
        ".svg": "image/svg+xml",
        ".pdf": "application/pdf",
        ".doc": "application/msword",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".xls": "application/vnd.ms-excel",
        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ".mp4": "video/mp4",
        ".mp3": "audio/mpeg",
        ".wav": "audio/wav",
        ".txt": "text/plain",
        ".json": "application/json",
        ".zip": "application/zip",
        ".html": "text/html",
        ".css": "text/css",
        ".js": "application/javascript",
    }
    return mapping.get(ext, "application/octet-stream")


def main():
    parser = argparse.ArgumentParser(description="迁移本地文件到京东云对象存储")
    parser.add_argument("--local-dir", default="/data/cfc-uploads", help="本地文件目录")
    parser.add_argument("--dry-run", action="store_true", help="仅预览，不实际上传")
    parser.add_argument("--overwrite", action="store_true", help="覆盖云端已存在的文件")
    args = parser.parse_args()

    local_dir = Path(args.local_dir)
    if not local_dir.is_dir():
        log.error("本地目录不存在: %s", local_dir)
        sys.exit(1)

    # 收集所有本地文件
    all_files = sorted(local_dir.rglob("*"))
    files = [f for f in all_files if f.is_file()]

    if not files:
        log.info("本地目录为空，无需迁移")
        return

    log.info("本地目录: %s", local_dir)
    log.info("目标 Bucket: %s", BUCKET)
    log.info("目标 Endpoint: %s", ENDPOINT)
    log.info("共发现 %d 个文件", len(files))
    if args.dry_run:
        log.info("模式: DRY-RUN（仅预览不上传）")
    elif args.overwrite:
        log.info("模式: 覆盖已存在文件")
    else:
        log.info("模式: 跳过已存在文件")

    # 连接京东云
    if not args.dry_run:
        log.info("正在连接京东云 OSS...")
        client = get_client()
        # 验证连通性
        try:
            client.head_bucket(Bucket=BUCKET)
            log.info("Bucket %s 连接成功", BUCKET)
        except Exception as e:
            log.error("Bucket 连接失败: %s", e)
            sys.exit(1)

        cloud_keys = list_cloud_keys(client) if not args.overwrite else set()
        log.info("云端已有 %d 个文件", len(cloud_keys))
    else:
        client = None
        cloud_keys = set()

    # 上传文件
    success = 0
    skipped = 0
    failed = 0

    for file_path in files:
        # 计算相对路径作为 key（保持目录结构）
        relative = file_path.relative_to(local_dir)
        key = "uploads/" + relative.as_posix()

        if client and key in cloud_keys and not args.overwrite:
            skipped += 1
            continue

        if upload_file(client, file_path, key, args.dry_run, args.overwrite, cloud_keys):
            success += 1
        else:
            failed += 1

    # 总结
    if not args.dry_run:
        log.info("=" * 50)
        log.info("迁移完成: 成功=%d, 跳过=%d, 失败=%d", success, skipped, failed)
        if failed > 0:
            log.warning("有 %d 个文件上传失败，请检查日志", failed)
    else:
        log.info("=" * 50)
        log.info("DRY-RUN 完成: 将上传 %d 个文件，跳过 %d 个（已存在）", success, skipped)


if __name__ == "__main__":
    main()