dickshow-index/download.py

"""Resumable MP3 downloader.

Downloads each episode's enclosure into audio/<epnum or guid>.mp3.
Skips files that already exist with the expected size (HTTP Range resume).
"""
import json
import os
import re
import sys
import time

import requests

from config import AUDIO_DIR, DOWNLOAD_TIMEOUT, CHUNK_SIZE
from feed import load_episodes


def slug(ep: dict) -> str:
    if ep.get("ep_num"):
        return f"ep{ep['ep_num']:04d}"
    g = re.sub(r"[^a-zA-Z0-9-]", "", ep["guid"])[:24]
    return f"bonus-{g}"


def audio_path(ep: dict) -> str:
    return os.path.join(AUDIO_DIR, slug(ep) + ".mp3")


def expected_size(ep: dict) -> int:
    return ep.get("audio_bytes") or 0


def download_episode(ep: dict, force: bool = False) -> str:
    """Download one episode. Returns local path. Raises on hard failure."""
    url = ep["audio_url"]
    if not url:
        raise ValueError(f"no audio_url for {ep['guid']}")
    dest = audio_path(ep)
    tmp = dest + ".part"
    want = expected_size(ep)

    if os.path.exists(dest):
        if force:
            os.remove(dest)
        else:
            have = os.path.getsize(dest)
            if want and have != want:
                print(f"  size mismatch for {dest} (have {have}, want {want}); "
                      f"redownloading", file=sys.stderr)
                os.remove(dest)
            else:
                return dest

    headers = {"User-Agent": "dickshow-indexer/1.0"}
    mode = "wb"
    if os.path.exists(tmp):
        have = os.path.getsize(tmp)
        if want and have == want:
            os.replace(tmp, dest)
            return dest
        if have > 0:
            headers["Range"] = f"bytes={have}-"
            mode = "ab"

    last_err = None
    for attempt in range(1, 6):
        try:
            with requests.get(url, headers=headers, stream=True,
                              timeout=DOWNLOAD_TIMEOUT, allow_redirects=True) as r:
                if r.status_code == 416:  # range beyond EOF -> already complete
                    os.replace(tmp, dest)
                    return dest
                r.raise_for_status()
                total = int(r.headers.get("Content-Length", "0") or 0)
                if mode == "ab" and r.status_code == 200:
                    # server ignored Range; restart cleanly
                    mode = "wb"
                with open(tmp, mode) as f:
                    for chunk in r.iter_content(chunk_size=CHUNK_SIZE):
                        if chunk:
                            f.write(chunk)
                have = os.path.getsize(tmp)
                if want and have < want:
                    raise IOError(f"incomplete download {have}/{want}")
                break
        except Exception as e:  # noqa: BLE001
            last_err = e
            print(f"  retry {attempt}: {e}", file=sys.stderr)
            time.sleep(5 * attempt)
            mode = "ab"
            headers["Range"] = f"bytes={os.path.getsize(tmp)}-" \
                if os.path.exists(tmp) and os.path.getsize(tmp) > 0 else {}
    else:
        raise IOError(f"download failed for {url}: {last_err}")

    os.replace(tmp, dest)
    return dest


def main() -> None:
    from config import TRANSCRIPT_DIR
    from transcribe import transcript_path
    eps = load_episodes()
    only = set()
    for a in sys.argv[1:]:
        only.add(int(a))
    ok = fail = skip = 0
    for ep in eps:
        if ep["ep_num"] is not None and only and ep["ep_num"] not in only:
            continue
        if os.path.exists(transcript_path(ep)):
            skip += 1          # already fully transcribed; audio not needed
            continue
        dest = audio_path(ep)
        if os.path.exists(dest) and not (expected_size(ep) and
                                         os.path.getsize(dest) != expected_size(ep)):
            skip += 1
            continue
        label = f"Ep {ep['ep_num']}" if ep["ep_num"] else ep["title"][:40]
        try:
            p = download_episode(ep)
            mb = os.path.getsize(p) / 1e6
            print(f"{label}: {mb:.1f} MB")
            ok += 1
        except Exception as e:  # noqa: BLE001
            print(f"{label}: FAILED - {e}", file=sys.stderr)
            fail += 1
        with open(os.path.join(AUDIO_DIR, ".progress.json"), "w") as f:
            json.dump({"last_done": ep["guid"]}, f)
    print(f"downloaded={ok} skipped={skip} failed={fail}")


if __name__ == "__main__":
    main()