"""
使用 cos-python-sdk-v5 全量扫描 0004/data/ 下的文件，
找出已有 SRT 文件的目录，将对应的 MP4 COS 路径写入 txt 文件。

安装依赖：pip install cos-python-sdk-v5
"""
import os
import re
from collections import defaultdict
from qcloud_cos import CosConfig, CosS3Client
from config.config_utils import config as app_config

SCAN_PREFIX = "0004/data/"
OUTPUT_TXT  = "./text/0004_has_srt.txt"
DIR_RE      = re.compile(r'.+/data/[^/]+/[^/]+/$')


def build_client() -> tuple[CosS3Client, str]:
    cos_cfg = app_config.get_cos_config()
    secret_id  = cos_cfg.get("secret_id")
    secret_key = cos_cfg.get("secret_key")
    region     = cos_cfg.get("region", "ap-chongqing")
    bucket     = cos_cfg.get("bucket", "asr-dig-data-1253181912")

    if not secret_id or not secret_key or "填入" in secret_id:
        raise ValueError("请先在 config/config.json 的 tencentcloud.cos 中填入 secret_id 和 secret_key")

    cfg    = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
    client = CosS3Client(cfg)
    return client, bucket


def list_all_objects(client: CosS3Client, bucket: str, prefix: str) -> list[str]:
    """分页获取 prefix 下全部 Key"""
    keys   = []
    marker = ""
    page   = 0

    while True:
        page += 1
        kwargs = dict(Bucket=bucket, Prefix=prefix, MaxKeys=1000)
        if marker:
            kwargs["Marker"] = marker

        resp     = client.list_objects(**kwargs)
        contents = resp.get("Contents", [])
        keys.extend(item["Key"] for item in contents)

        print(f"\r  第 {page} 页，已获取 {len(keys)} 个文件...", end="", flush=True)

        if resp.get("IsTruncated") == "true":
            marker = resp.get("NextMarker", "")
        else:
            break

    print()
    return keys


def main():
    client, bucket = build_client()
    bucket_prefix  = f"cos://{bucket}"

    print(f"正在全量扫描 {bucket_prefix}/{SCAN_PREFIX} ...")
    keys = list_all_objects(client, bucket, SCAN_PREFIX)
    print(f"共找到 {len(keys)} 个文件")

    if not keys:
        print("未找到任何文件，请检查密钥和路径是否正确")
        return

    # 按父目录分组
    dir_files: dict[str, list[str]] = defaultdict(list)
    for key in keys:
        idx      = key.rfind("/")
        dir_path = key[:idx + 1]   # e.g. "0004/data/BV1xx/BV1xx/"
        filename = key[idx + 1:]
        if filename:
            dir_files[dir_path].append(filename)

    # 找有 SRT 的目录，收集 MP4 路径
    mp4_paths   = []
    has_srt_cnt = 0
    no_srt_cnt  = 0

    for dir_path, filenames in sorted(dir_files.items()):
        # 只处理 data/{id}/{id}/ 结构
        if not DIR_RE.match(dir_path):
            continue

        has_srt   = any(f.lower().endswith(".srt") for f in filenames)
        mp4_files = [f for f in filenames if f.lower().endswith(".mp4")]

        if has_srt:
            has_srt_cnt += 1
            for mp4 in mp4_files:
                mp4_paths.append(f"{bucket_prefix}/{dir_path}{mp4}")
        else:
            no_srt_cnt += 1

    print(f"有 SRT 的目录: {has_srt_cnt} 个  →  {len(mp4_paths)} 个 MP4")
    print(f"无 SRT 的目录: {no_srt_cnt} 个")

    if not mp4_paths:
        print("没有找到含 SRT 的目录")
        return

    os.makedirs("./text", exist_ok=True)
    with open(OUTPUT_TXT, "w", encoding="utf-8") as f:
        for path in mp4_paths:
            f.write(path + "\n")

    print(f"已写入: {OUTPUT_TXT}  ({len(mp4_paths)} 条)")


if __name__ == "__main__":
    main()