#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""批量提取音频指纹并上传到 ACRCloud 自定义曲库。

流程：
1. 扫描目标目录的音频文件
2. 用 acrcloud_extr_mac 提取 DB 指纹（.bin 文件）
3. POST 到 ACRCloud 管理 API，data_type=fingerprint
4. 记录上传状态到 --state-file，支持断点续传

用法:
    python scripts/upload_acrcloud.py \
        --bucket-id 2336 \
        --audio-dir /Volumes/移动硬盘/composition_test \
        --ext .wav

    # 断点续传（跳过已上传成功的文件）：
    python scripts/upload_acrcloud.py \
        --bucket-id YOUR_BUCKET_ID \
        --audio-dir /Volumes/移动硬盘/composition_test \
        --state-file upload_state.json
"""

import argparse
import json
import logging
import os
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

logger = logging.getLogger(__name__)

SUPPORTED_EXTENSIONS = {".mp3", ".wav", ".flac", ".ogg", ".m4a", ".aac", ".wma", ".flac"}
EXTR_TOOL = Path(__file__).resolve().parent / "acrcloud_extr_tools-master" / "acrcloud_extr_mac"
API_BASE = "https://cn-api-v2.acrcloud.cn/api"


# ===== 加载 .env =====

def _load_env_file() -> None:
    env_path = Path(__file__).resolve().parent.parent.parent / ".env"
    if not env_path.exists():
        return
    with env_path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            if not key:
                continue
            os.environ.setdefault(key, value.strip().strip('"').strip("'"))


_load_env_file()


# ===== 文件发现 =====

def _extract_song_id(path: str) -> str:
    """从路径中提取 song_id：取文件名第一段（下划线前），若为纯数字则使用，否则用路径哈希。"""
    import hashlib
    name = Path(path).stem
    prefix = name.split("_")[0]
    if prefix.isdigit():
        return prefix
    return str(int(hashlib.md5(path.encode()).hexdigest()[:8], 16))


def discover_audio_files(audio_dir: str | None, file_list: str | None, ext: str) -> list[tuple[str, str]]:
    """返回 [(song_id, audio_path), ...] 列表。"""
    results = []

    if file_list:
        with open(file_list, "r", encoding="utf-8") as f:
            for line in f:
                path = line.strip()
                if path:
                    results.append((_extract_song_id(path), path))
    elif audio_dir:
        for audio_file in sorted(Path(audio_dir).rglob(f"*{ext}")):
            if audio_file.is_file() and not audio_file.name.startswith("._"):
                results.append((_extract_song_id(str(audio_file)), str(audio_file)))
    else:
        logger.error("请指定 --audio-dir 或 --file-list")
        sys.exit(1)

    return results


# ===== 指纹提取 =====

def extract_fingerprint(audio_path: str, output_dir: str) -> Path | None:
    """调用 acrcloud_extr_mac 提取 DB 指纹，返回生成的 .bin 文件路径，失败返回 None。

    不加 -cli 表示生成 DB 指纹（用于入库），加 -cli 才是查询指纹。
    -o 必须传完整文件路径（含文件名），不能传目录。
    """
    if not EXTR_TOOL.exists():
        raise FileNotFoundError(f"acrcloud_extr_mac 不存在: {EXTR_TOOL}")

    out_file = Path(output_dir) / (Path(audio_path).stem + ".bin")
    cmd = [str(EXTR_TOOL), "-i", audio_path, "-o", str(out_file)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        logger.warning("指纹提取失败: %s\nstderr: %s", audio_path, result.stderr.strip())
        return None

    if not out_file.exists() or out_file.stat().st_size == 0:
        logger.warning("指纹文件不存在或为空: %s", out_file)
        return None

    return out_file


# ===== ACRCloud 上传 =====


def _encode_multipart(fields: dict[str, str], files: dict[str, tuple[str, bytes]]) -> tuple[str, bytes]:
    """构造 multipart/form-data 请求体。

    files: {field_name: (filename, content_bytes)}
    """
    boundary = f"----ACRCloudUpload{int(time.time())}"
    CRLF = b"\r\n"
    body = b""

    for key, value in fields.items():
        body += f"--{boundary}\r\n".encode()
        body += f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode()
        body += value.encode() + CRLF

    for key, (filename, content) in files.items():
        body += f"--{boundary}\r\n".encode()
        body += f'Content-Disposition: form-data; name="{key}"; filename="{filename}"\r\n'.encode()
        body += b"Content-Type: application/octet-stream\r\n\r\n"
        body += content + CRLF

    body += f"--{boundary}--\r\n".encode()
    content_type = f"multipart/form-data; boundary={boundary}"
    return content_type, body


def upload_fingerprint(
    bucket_id: str,
    fp_path: Path,
    title: str,
    access_token: str,
    user_defined: dict | None = None,
    timeout: int = 30,
) -> tuple[bool, str]:
    """上传指纹文件到 ACRCloud。

    Returns:
        (success, message)
    """
    url = f"{API_BASE}/buckets/{bucket_id}/files"

    fp_bytes = fp_path.read_bytes()

    fields: dict[str, str] = {
        "data_type": "fingerprint",
        "title": title,
    }
    if user_defined:
        fields["user_defined"] = json.dumps(user_defined, ensure_ascii=False)

    files = {"file": (fp_path.name, fp_bytes)}
    content_type, body = _encode_multipart(fields, files)

    req = urllib.request.Request(url, data=body, method="POST")
    req.add_header("Authorization", f"Bearer {access_token}")
    req.add_header("Content-Type", content_type)

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            return True, raw
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8", errors="replace")
        return False, f"HTTP {e.code}: {raw}"
    except Exception as e:
        return False, str(e)


# ===== 状态文件（断点续传）=====

def load_state(state_file: str) -> dict[str, str]:
    """加载上传状态，返回 {audio_path: 'ok'|'failed'} 字典。"""
    if not Path(state_file).exists():
        return {}
    with open(state_file, encoding="utf-8") as f:
        return json.load(f)


def save_state(state_file: str, state: dict[str, str]) -> None:
    with open(state_file, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


# ===== 主流程 =====

def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

    parser = argparse.ArgumentParser(description="批量提取音频指纹并上传到 ACRCloud 自定义曲库")
    parser.add_argument("--bucket-id", required=True, help="ACRCloud bucket ID")
    parser.add_argument("--audio-dir", help="音频文件目录")
    parser.add_argument("--file-list", help="音频文件路径列表文件（每行一个路径）")
    parser.add_argument("--ext", default=".wav", help="音频文件扩展名（默认 .wav）")
    parser.add_argument("--state-file", default="upload_acrcloud_state.json", help="上传状态文件，用于断点续传（默认 upload_acrcloud_state.json）")
    parser.add_argument("--retry-failed", action="store_true", help="重试上次失败的文件")
    parser.add_argument("--dry-run", action="store_true", help="只列出待上传文件，不实际上传")
    args = parser.parse_args()

    access_token = os.environ["ACRCLOUD_ACCESS_TOKEN"]

    # 检查 extr 工具
    if not EXTR_TOOL.exists():
        logger.error("找不到 acrcloud_extr_mac: %s", EXTR_TOOL)
        sys.exit(1)
    if not os.access(str(EXTR_TOOL), os.X_OK):
        logger.info("acrcloud_extr_mac 无执行权限，尝试 chmod +x")
        os.chmod(str(EXTR_TOOL), 0o755)

    audio_files = discover_audio_files(args.audio_dir, args.file_list, args.ext)
    logger.info("发现 %d 个音频文件", len(audio_files))

    state = load_state(args.state_file)

    # 过滤已成功上传的文件
    pending = []
    for song_id, audio_path in audio_files:
        status = state.get(audio_path)
        if status == "ok":
            continue
        if status == "failed" and not args.retry_failed:
            continue
        pending.append((song_id, audio_path))

    logger.info("待上传: %d 个（已跳过 %d 个成功 + %d 个失败）",
                len(pending),
                sum(1 for s in state.values() if s == "ok"),
                sum(1 for s in state.values() if s == "failed") if not args.retry_failed else 0)

    if args.dry_run:
        for song_id, audio_path in pending:
            print(f"  [dry-run] {audio_path} (song_id={song_id})")
        return

    success_count = 0
    fail_count = 0

    with tempfile.TemporaryDirectory() as tmp_dir:
        for i, (song_id, audio_path) in enumerate(pending, 1):
            logger.info("[%d/%d] 处理: %s", i, len(pending), Path(audio_path).name)

            # 1. 提取指纹
            t0 = time.perf_counter()
            fp_path = extract_fingerprint(audio_path, tmp_dir)
            fp_ms = (time.perf_counter() - t0) * 1000

            if fp_path is None:
                logger.error("  指纹提取失败，跳过")
                state[audio_path] = "failed"
                fail_count += 1
                save_state(args.state_file, state)
                continue

            logger.info("  指纹提取完成: %.0fms, 大小=%d bytes", fp_ms, fp_path.stat().st_size)

            # 2. 上传
            title = Path(audio_path).stem      # 如 203648006_GA000003
            user_defined = {"song_id": song_id}

            t1 = time.perf_counter()
            success, msg = upload_fingerprint(
                bucket_id=args.bucket_id,
                fp_path=fp_path,
                title=title,
                access_token=access_token,
                user_defined=user_defined,
            )
            upload_ms = (time.perf_counter() - t1) * 1000

            if success:
                logger.info("  上传成功: %.0fms | %s", upload_ms, msg[:120])
                state[audio_path] = "ok"
                success_count += 1
            else:
                logger.error("  上传失败: %.0fms | %s", upload_ms, msg)
                state[audio_path] = "failed"
                fail_count += 1

            # 每条都保存状态，确保中断后可续传
            save_state(args.state_file, state)

            # 清理临时指纹文件，避免 tmp 目录膨胀
            fp_path.unlink(missing_ok=True)

    logger.info("上传完成: 成功 %d, 失败 %d，状态已保存到 %s",
                success_count, fail_count, args.state_file)


if __name__ == "__main__":
    main()