"""批量下载：并发、可断点续跑、单条失败不影响整批。

设计要点：
  - 设备身份只注册一次，各 worker 共用；但每个 worker 有独立 Session
    （requests.Session 并发下不保证安全）
  - 落盘文件名就是 aweme_id，所以「是否已下载」直接看文件在不在，
    不依赖去重记录 —— 记录丢了也能无损续跑
  - 失败条目写 failed.txt，跑完可直接拿它当输入重试
"""

from __future__ import annotations

import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Sequence

import requests

from .app import AppClient, Device, NotAVideo
from .client import DouyinError
from .download import Registry, download


class DevicePool:
    """设备身份池。

    实测单个 device_id 在约 575 次 detail 调用后被限流，表现是接口一直
    回 HTTP 200 空 body；同一 IP 下换个新注册的设备立刻恢复 —— 说明配额
    是设备级而非 IP 级。注册本身免签名、成本极低，所以按量轮换即可。

    用 generation 标记当前设备：worker 拿到的 gen 和池里的对得上才允许
    触发轮换，避免多个 worker 同时发现失败、把设备连换好几次。
    """

    def __init__(self, max_calls: int = 400, initial: Optional[Device] = None) -> None:
        self._lock = threading.Lock()
        self._device = initial
        self._gen = 0
        self._calls = 0
        self.max_calls = max_calls
        self.rotations = 0

    def current(self) -> tuple[Device, int]:
        with self._lock:
            if self._device is None or self._calls >= self.max_calls:
                self._rotate_locked("配额用尽")
            self._calls += 1
            return self._device, self._gen

    def report_dead(self, gen: int) -> None:
        """worker 发现当前设备被限流时调用。"""
        with self._lock:
            if gen == self._gen:
                self._rotate_locked("被限流")

    def _rotate_locked(self, why: str) -> None:
        self._device = Device.register()
        self._gen += 1
        self._calls = 0
        self.rotations += 1
        print(f"  [设备轮换#{self.rotations}] {why} -> "
              f"device_id={self._device.device_id}", flush=True)


@dataclass
class BatchStats:
    total: int = 0
    done: int = 0
    skipped: int = 0
    not_video: int = 0
    failed: int = 0
    bytes_: int = 0
    lock: threading.Lock = field(default_factory=threading.Lock)

    @property
    def finished(self) -> int:
        return self.done + self.skipped + self.not_video + self.failed


def run_batch(aweme_ids: Sequence[str], out_dir: Path, device: Optional[Device] = None,
              workers: int = 5, delay: float = 0.3,
              registry: Optional[Registry] = None,
              log_every: int = 25, max_calls_per_device: int = 400) -> BatchStats:
    out_dir.mkdir(parents=True, exist_ok=True)
    registry = registry or Registry(out_dir / ".downloaded.json")
    pool = DevicePool(max_calls=max_calls_per_device, initial=device)

    stats = BatchStats(total=len(aweme_ids))
    failed_path = out_dir / "failed.txt"
    # 每轮重来，否则上一轮已修复的失败会一直留在文件里造成误导
    failed_path.write_text("", encoding="utf-8")
    failed_lock = threading.Lock()
    local = threading.local()
    t0 = time.time()

    def client_for(device_: Device, gen: int) -> AppClient:
        """每线程每代设备一个 Session（Session 并发不安全）。"""
        if getattr(local, "gen", None) != gen:
            local.client = AppClient(device=device_)
            local.gen = gen
        return local.client

    def dl_session() -> requests.Session:
        """下载用的 Session：CDN 免鉴权，不该占设备配额，也不随设备轮换重建。"""
        if not hasattr(local, "dl"):
            local.dl = requests.Session()
            local.dl.headers["User-Agent"] = (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                "(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
            )
        return local.dl

    def record_failure(aweme_id: str, why: str) -> None:
        with failed_lock:
            with open(failed_path, "a", encoding="utf-8") as f:
                f.write(f"{aweme_id}\t{why}\n")

    def fetch_with_rotation(aweme_id: str):
        """限流表现为空 body。碰到就换设备再试一次，而不是直接判失败。"""
        last: Optional[Exception] = None
        for _ in range(2):
            dev, gen = pool.current()
            client = client_for(dev, gen)
            if delay:
                time.sleep(delay)
            try:
                return client.fetch_detail(aweme_id, retries=2)
            except NotAVideo:
                raise
            except DouyinError as exc:
                last = exc
                if "空 body" in str(exc):
                    pool.report_dead(gen)
                    continue
                raise
        raise last  # type: ignore[misc]

    def handle(aweme_id: str) -> None:
        dest = out_dir / f"{aweme_id}.mp4"
        # 文件在就算下过，不必碰网络也不依赖记录
        if dest.exists() and dest.stat().st_size > 0:
            with stats.lock:
                stats.skipped += 1
            return

        try:
            aweme = fetch_with_rotation(aweme_id)
        except NotAVideo as exc:
            with stats.lock:
                stats.not_video += 1
            record_failure(aweme_id, f"not_video: {exc}")
            return
        except Exception as exc:
            with stats.lock:
                stats.failed += 1
            record_failure(aweme_id, f"detail: {type(exc).__name__}: {exc}")
            return

        if not aweme.video_urls:
            with stats.lock:
                stats.not_video += 1
            record_failure(aweme_id, "no_video_url")
            return

        try:
            result = download(dl_session(), aweme.video_urls, dest, progress=False)
        except Exception as exc:
            with stats.lock:
                stats.failed += 1
            record_failure(aweme_id, f"download: {type(exc).__name__}: {exc}")
            return

        registry.add(aweme_id, result.path, result.size,
                     desc=aweme.desc, author=aweme.author, uri=aweme.uri)
        with stats.lock:
            stats.done += 1
            stats.bytes_ += result.size

    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = [executor.submit(handle, a) for a in aweme_ids]
        for n, _ in enumerate(as_completed(futures), 1):
            if n % log_every == 0 or n == len(futures):
                el = time.time() - t0
                rate = n / el if el else 0
                eta = (len(futures) - n) / rate if rate else 0
                print(f"[{n}/{len(futures)}] 成功{stats.done} 跳过{stats.skipped} "
                      f"图集{stats.not_video} 失败{stats.failed} "
                      f"| {stats.bytes_ / 1024 ** 3:.2f}GB "
                      f"| {rate * 60:.0f}条/分 | 换设备{pool.rotations}次 "
                      f"| 剩约 {eta / 60:.0f}分",
                      flush=True)

    registry.flush()
    return stats