#!/usr/bin/env python3
"""Re-download the pdfs/ files that are actually Drive interstitial pages.

A failed download is detected by content, not size: a real PDF starts with
the magic bytes ``%PDF``. Anything else (e.g. the "Google Drive can't scan
this file for viruses" interstitial HTML) is re-downloaded using the full
confirm+uuid flow with cookie persistence.

Only replaces the local file when the new bytes look like a real PDF.
Prints FAIL lines for files that could not be fetched (these go into
manual_download.html).

Usage:
    ./retry_download.py [--dry-run]
"""
import csv
import re
import sys
import urllib.request
from http.cookiejar import CookieJar
from pathlib import Path

HERE = Path(__file__).resolve().parent
PDFS = HERE / "pdfs"

UA = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)"}


def drive_id(link):
    m = re.search(r"/file/d/([^/?]+)", link)
    return m.group(1) if m else None


def fetch(url, opener):
    req = urllib.request.Request(url, headers=UA)
    with opener.open(req, timeout=120) as r:
        return r.read()


def download_drive_file(file_id, dest):
    """Return True and write dest only if we got PDF bytes."""
    opener = urllib.request.build_opener(
        urllib.request.HTTPCookieProcessor(CookieJar())
    )
    page = fetch(
        f"https://drive.google.com/uc?export=download&id={file_id}", opener
    )
    if page.startswith(b"%PDF"):
        data = page
    else:
        text = page.decode("utf-8", "replace")
        # Modern interstitial: <form action="https://drive.usercontent...">
        # with hidden inputs id/export/confirm/uuid.
        action = re.search(
            r'<form[^>]*action="([^"]+)"', text
        )
        fields = dict(
            re.findall(r'name="(\w+)" value="([^"]*)"', text)
        )
        if not action or "uuid" not in fields:
            return False, "no uuid/confirm token in interstitial"
        query = "&".join(f"{k}={v}" for k, v in fields.items())
        data = fetch(f"{action.group(1)}?{query}", opener)
    if not data.startswith(b"%PDF"):
        return False, f"got {len(data)} non-PDF bytes"
    if len(data) < 50_000:
        return False, f"suspiciously small PDF ({len(data)} bytes)"
    dest.write_bytes(data)
    return True, f"{len(data) / 1e6:.1f} MB"


def main():
    dry = "--dry-run" in sys.argv
    with (HERE / "download.txt").open(encoding="utf-8-sig") as f:
        rows = list(csv.DictReader(f, delimiter="\t"))

    # Mirror download_all.sh naming for duplicate titles.
    seen: dict = {}
    targets = []
    for row in rows:
        name = (row.get("Book Name") or "").strip()
        if not name:
            continue
        n = seen.get(name, 0)
        seen[name] = n + 1
        target = name
        if n:
            p = Path(name)
            target = f"{p.stem} ({n + 1}){p.suffix}" if p.suffix else f"{name} ({n + 1})"
        targets.append((target, row.get("Link") or ""))

    failed = []
    for target, link in targets:
        path = PDFS / target
        if path.is_file() and path.read_bytes()[:4] == b"%PDF":
            continue  # genuine PDF, leave it alone
        fid = drive_id(link)
        if not fid:
            failed.append((target, link, "no file id"))
            print(f"FAIL {target}: no file id")
            continue
        if dry:
            size_info = f"{path.stat().st_size} bytes" if path.is_file() else "missing"
            print(f"WOULD-RETRY {target} ({size_info})")
            continue
        print(f"DOWNLOAD {target} ...", flush=True)
        try:
            ok, msg = download_drive_file(fid, path)
        except Exception as e:  # noqa: BLE001 - report, don't stop the batch
            ok, msg = False, f"{type(e).__name__}: {e}"
        print(f"  {'OK' if ok else 'FAIL'}: {msg}")
        if not ok:
            failed.append((target, link, msg))

    if failed and not dry:
        lines = [
            "<!DOCTYPE html><html><head><meta charset='utf-8'>"
            "<title>Manual downloads</title></head><body>",
            "<h1>Manual downloads</h1>",
            "<p>These files could not be downloaded automatically. "
            "Open each link and download it by hand into <code>pdfs/</code>.</p>",
            "<ul>",
        ]
        for target, link, msg in failed:
            import html as _h

            lines.append(
                f"<li><a href='{_h.escape(link)}'>{_h.escape(target)}</a>"
                f" — auto-download failed: {_h.escape(msg)}</li>"
            )
        lines += ["</ul></body></html>"]
        (HERE / "manual_download.html").write_text(
            "\n".join(lines), encoding="utf-8"
        )
        print(f"Wrote manual_download.html with {len(failed)} links.")
    elif not dry:
        (HERE / "manual_download.html").unlink(missing_ok=True)
        print("All suspect files re-downloaded successfully.")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())