#!/usr/bin/env python3
"""
Collect every episode and movie URL from History Vault, A&E Crime Central and
play.history.com.
No browser or login needed: the show/movie lists are in the pages' server-rendered data.

For each site, writes (next to this script, or into the folder given with --out):
  <site>_episode_urls.txt   one full episode URL per line
  <site>_episodes.csv       series, season, episode, title, url, needs_login
  <site>_movie_urls.txt     one full movie URL per line (only if the site has movies)

Run: python3 extract-svod-urls.py [historyvault] [aecrimecentral] [playhistory] [lifetimemovieclub] [--out DIR]
     (no site names = all except lifetimemovieclub, which needs ~/.aetn_cookies.json)
"""
import csv, json, os, re, sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from html import unescape

SITES = {
    "historyvault":   "https://watch.historyvault.com",
    "aecrimecentral": "https://watch.aecrimecentral.com",
    "playhistory":    "https://play.history.com",
    # Needs a login: pages return 404 unless the cookie file below is available
    "lifetimemovieclub": "https://watch.lifetimemovieclub.com",
}

HERE = os.path.dirname(os.path.abspath(__file__))
COOKIE_FILES = [os.path.expanduser("~/.aetn_cookies.json"),
                os.path.join(HERE, "..", "data", ".aetn_cookies.json")]


@lru_cache(maxsize=None)
def cookie_header(base):
    """Cookie header built from the cookie file saved by save_cookies.sh, if there is one."""
    host = base.split("//")[1]
    for path in COOKIE_FILES:
        try:
            with open(path) as f:
                cookies = json.load(f)
        except (OSError, ValueError):
            continue
        return "; ".join(f"{c['name']}={c['value']}" for c in cookies
                         if c.get("domain") and host.endswith(c["domain"].lstrip(".")))
    return ""


def apollo_state(base, path):
    headers = {"User-Agent": "Mozilla/5.0"}
    if cookie_header(base):
        headers["Cookie"] = cookie_header(base)
    req = urllib.request.Request(base + path, headers=headers)
    with urllib.request.urlopen(req, timeout=60) as resp:
        html = resp.read().decode("utf-8", "replace")
    m = re.search(r'<script id="__NEXT_DATA__"[^>]*>(.*?)</script>', html, re.S)
    return json.loads(m.group(1))["props"].get("apolloState", {})


def canonicals(base, list_path, prefix):
    try:
        state = apollo_state(base, list_path)
    except Exception:
        return []
    return sorted({v["canonical"] for k, v in state.items()
                   if k.startswith(prefix) and v.get("canonical")})


def episodes_of(base, path):
    for attempt in range(3):
        try:
            state = apollo_state(base, path)
            break
        except Exception as e:
            if attempt == 2:
                print(f"  FAILED {path}: {e}", file=sys.stderr)
                return []
    series = next((v for k, v in state.items()
                   if k.startswith("Series:") and v.get("canonical") == path), {})
    rows = []
    for key, ep in state.items():
        # only this show's episodes; the page data also carries recommended shows
        if not key.startswith("Episode:") or not (ep.get("canonical") or "").startswith(path + "/"):
            continue
        video = state.get((ep.get("primaryVideo") or {}).get("__ref", ""), {})
        rows.append({
            "series": unescape(series.get("title") or path.rsplit("/", 1)[-1]),
            "season": ep.get("tvSeasonNumber"),
            "episode": ep.get("tvSeasonEpisodeNumber"),
            "title": unescape(ep.get("title") or ""),
            "url": base + ep["canonical"],
            "needs_login": bool(video.get("isBehindWall")),
        })
    return rows


def scrape(site, out_dir):
    base = SITES[site]
    paths = canonicals(base, "/shows", "Series:")
    print(f"[{site}] {len(paths)} series found, reading episode lists...")
    with ThreadPoolExecutor(max_workers=4) as pool:
        rows = [r for batch in pool.map(lambda p: episodes_of(base, p), paths) for r in batch]
    rows.sort(key=lambda r: (r["series"], r["season"] or 0, r["episode"] or 0))

    if rows:
        with open(os.path.join(out_dir, f"{site}_episode_urls.txt"), "w") as f:
            f.write("\n".join(r["url"] for r in rows) + "\n")
        with open(os.path.join(out_dir, f"{site}_episodes.csv"), "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(rows[0]))
            w.writeheader()
            w.writerows(rows)
    locked = sum(r["needs_login"] for r in rows)
    print(f"[{site}] saved {len(rows)} episode URLs ({locked} need a login)")

    # Standalone documentaries and specials are movie-like: /specials on some sites, and only
    # linked from the /shows and home pages on others
    specials = {p for page in ("/specials", "/shows", "/")
                for p in canonicals(base, page, "Special:") if p.startswith("/specials/")}
    movies = sorted(set(canonicals(base, "/movies", "Movie:")) | specials)
    if movies:
        with open(os.path.join(out_dir, f"{site}_movie_urls.txt"), "w") as f:
            f.write("\n".join(base + p for p in movies) + "\n")
    print(f"[{site}] saved {len(movies)} movie URLs ({len(specials)} of them specials)")


def main():
    args = sys.argv[1:]
    out_dir = os.path.dirname(os.path.abspath(__file__))
    if "--out" in args:
        i = args.index("--out")
        out_dir = args[i + 1]
        del args[i:i + 2]
    unknown = [a for a in args if a not in SITES]
    if unknown:
        sys.exit(f"Unknown site(s): {', '.join(unknown)}. Choose from: {', '.join(SITES)}")
    for site in args or [s for s in SITES if s != "lifetimemovieclub"]:
        scrape(site, out_dir)


if __name__ == "__main__":
    main()