#!/usr/bin/env python3
"""Resolve Letterboxd poster image URLs for every film in a ratings.csv export.

Usage:
    python fetch_posters.py ratings.csv [-o posters.json]

Writes a {"https://boxd.it/xxxx": "https://a.ltrbxd.com/..."} map. The output file
doubles as the cache: re-running after a new export only fetches the new rows.
"""

import argparse
import csv
import json
import os
import re
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor

POSTER_RE = re.compile(r'https://a\.ltrbxd\.com/resized/film-poster/[^"\'\ <>]+')

# Some films serve their poster from /resized/sm/upload/ instead. Those pages also carry
# landscape stills, so match the poster crop shape and keep only the 2:3 images.
FALLBACK_RE = re.compile(r'https://a\.ltrbxd\.com/resized/[^"\'\ <>]*?-0-(\d+)-0-(\d+)-crop\.jpg(?:\?v=[0-9a-z]+)?')

# Letterboxd rejects the default Python-urllib agent with a 403.
USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/126.0 Safari/537.36"
)

URI_COLUMN = "Letterboxd URI"
WORKERS = 4
REQUEST_DELAY = 0.15
TIMEOUT = 20
SAVE_EVERY = 20

print_lock = threading.Lock()
cache_lock = threading.Lock()


def read_uris(csv_path):
    with open(csv_path, "r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        if reader.fieldnames is None or URI_COLUMN not in reader.fieldnames:
            sys.exit(
                "'%s' has no '%s' column. Columns found: %s"
                % (csv_path, URI_COLUMN, reader.fieldnames)
            )
        seen, uris = set(), []
        for row in reader:
            uri = (row.get(URI_COLUMN) or "").strip()
            if uri and uri not in seen:
                seen.add(uri)
                uris.append(uri)
    return uris


def load_cache(path):
    if not os.path.exists(path):
        return {}
    try:
        with open(path, "r", encoding="utf-8") as handle:
            data = json.load(handle)
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError) as error:
        print("Could not read cache %s (%s), starting fresh." % (path, error))
        return {}


def save_cache(path, cache):
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as handle:
        json.dump(cache, handle, indent=1, sort_keys=True, ensure_ascii=False)
    os.replace(tmp, path)


def find_poster(body):
    match = POSTER_RE.search(body)
    if match:
        return match.group(0)
    for match in FALLBACK_RE.finditer(body):
        width, height = int(match.group(1)), int(match.group(2))
        if width > 0 and abs(float(height) / width - 1.5) < 0.2:
            return match.group(0)
    return None


def fetch_poster(uri):
    request = urllib.request.Request(uri, headers={"User-Agent": USER_AGENT})
    for attempt in range(2):
        try:
            with urllib.request.urlopen(request, timeout=TIMEOUT) as response:
                body = response.read().decode("utf-8", "replace")
            return find_poster(body)
        except (urllib.error.URLError, OSError) as error:
            if attempt == 1:
                with print_lock:
                    print("  failed %s: %s" % (uri, error))
                return None
            time.sleep(1.0)
    return None


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("csv_path", help="Letterboxd ratings.csv export")
    parser.add_argument("-o", "--output", default=None, help="output json (default: posters.json next to the csv)")
    parser.add_argument("--retry-failed", action="store_true", help="also retry URIs that previously had no poster")
    args = parser.parse_args()

    output = args.output or os.path.join(os.path.dirname(os.path.abspath(args.csv_path)), "posters.json")

    uris = read_uris(args.csv_path)
    cache = load_cache(output)
    if args.retry_failed:
        cache = {key: value for key, value in cache.items() if value}

    pending = [uri for uri in uris if uri not in cache]
    print("%d films, %d already cached, %d to fetch." % (len(uris), len(uris) - len(pending), len(pending)))
    if not pending:
        print("Nothing to do. %s is up to date." % output)
        return

    done = [0]

    def worker(uri):
        time.sleep(REQUEST_DELAY)
        poster = fetch_poster(uri)
        with cache_lock:
            cache[uri] = poster
            done[0] += 1
            count = done[0]
            if count % SAVE_EVERY == 0:
                save_cache(output, cache)
        with print_lock:
            print("  %d/%d" % (count, len(pending)), end="\r", flush=True)

    try:
        with ThreadPoolExecutor(max_workers=WORKERS) as pool:
            list(pool.map(worker, pending))
    except KeyboardInterrupt:
        print("\nInterrupted, saving what we have.")
    finally:
        with cache_lock:
            save_cache(output, cache)

    found = sum(1 for uri in uris if cache.get(uri))
    missing = [uri for uri in uris if not cache.get(uri)]
    print("\nWrote %s: %d of %d films have a poster." % (output, found, len(uris)))
    if missing:
        print("%d without a poster (they show as text cards): %s" % (len(missing), ", ".join(missing[:5])))
        print("Re-run with --retry-failed to try them again.")


if __name__ == "__main__":
    main()