#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import csv
import subprocess
from pathlib import Path
from urllib.parse import unquote, urlparse


def main():
    csv_path = Path("/hdd1b/cryosparc/nbdbStatic/script/emdb_entry.csv")
    base_dir = Path("/hdd1b/cryosparc/nbdbStatic/EMDB")
    log_path = Path("/hdd1b/cryosparc/nbdbStatic/script/download-failed.txt")
    max_rows = 60500  # 从第2行开始，最多处理x行

    with csv_path.open("r", encoding="utf-8-sig", newline="") as csv_file:
        reader = csv.reader(csv_file, delimiter=";")
        next(reader, None)  # 跳过第1行表头
        processed = 0
        failure_records = []

        for row in reader:
            if processed >= max_rows:
                break
            if not row or not str(row[0]).strip():
                continue

            first_col = str(row[0]).strip()
            xxxxx = first_col
            url = str(row[2]).strip() if len(row) > 1 else ""

            if not url:
                print(f"跳过（第二列没有下载链接）：{xxxxx}")
                failure_records.append((xxxxx, "未匹配"))
                processed += 1
                continue

            target_dir = base_dir / xxxxx
            original_name = Path(unquote(urlparse(url).path)).name

            if not original_name:
                print(f"跳过（无法从链接取得文件名）：{xxxxx}，{url}")
                failure_records.append((xxxxx, "未匹配"))
                processed += 1
                continue

            target_file = target_dir / original_name

            if target_file.exists():
                print(f"已存在，跳过下载：{target_file}")
                processed += 1
                continue

            target_dir.mkdir(parents=True, exist_ok=True)

            cmd = [
                "wget",
                "-P",
                str(target_dir),
                url,
            ]

            print("执行:", " ".join(cmd))
            try:
                result = subprocess.run(cmd, check=False)
                if result.returncode != 0:
                    failure_records.append((xxxxx, "下载失败"))
                    print(f"下载失败：ID={xxxxx}，错误码={result.returncode}，继续下一条")
            except OSError as error:
                failure_records.append((xxxxx, "下载失败"))
                print(f"下载命令执行失败：ID={xxxxx}，{error}，继续下一条")

            processed += 1

        log_path.parent.mkdir(parents=True, exist_ok=True)
        unique_records = list(dict.fromkeys(failure_records))
        log_path.write_text(
            "\n".join(f"{entry_id},{status}" for entry_id, status in unique_records)
            + ("\n" if unique_records else ""),
            encoding="utf-8",
        )
        unmatched_count = sum(status == "未匹配" for _, status in unique_records)
        download_failed_count = sum(status == "下载失败" for _, status in unique_records)
        print(
            f"处理完成，未匹配 {unmatched_count} 个，"
            f"下载失败 {download_failed_count} 个，日志：{log_path}"
        )


if __name__ == "__main__":
    main()