#!/usr/bin/env python3
"""
Find first-publication years for the Project Gutenberg works in overview/data.json.

Usage:
    python years.py                 # after build_overview.py
    python years.py --cached-only   # no network, use years_cache.json only

Gutenberg's catalogue has no publication dates, so the page used to place its works by the
author's birth decade. This adds a year from two sources, in order:

1. Wikidata: items carrying the Gutenberg ebook ID (P2034). The year of the work the ebook is
   an edition or translation of (P629 -> P577) is preferred, else the item's own date (P577).
   The work's Wikidata ID is kept (field "wk") so editions and translations can be grouped.
2. Open Library: a title + author-surname search; the earliest year of the best-matching
   record whose title is close to the Gutenberg title.
3. The ebook's own title page: the "Original publication:" line that newer Gutenberg ebooks
   carry, else a year on the transcribed title page (publisher lines, "Copyright, 1887", a year
   on its own line). The first 60 KB of each text come from Gutenberg's mirror
   gutenberg.pglaf.org and are cached in cache/pgheads/.
4. Otherwise, if the author's birth and death years are known, the midpoint of the author's
   life, marked as an estimate ("ys": "lifespan").

A year is only accepted if it is plausible for the first author (from age 12 to 60 years after
death, or within 80 years of birth when no death date is known) and before 1971, when Project
Gutenberg began, which drops ebook release dates and most modern reprints. Writes
"y" (year), "ys" ("wikidata", "openlibrary", "titlepage" or "lifespan") and, where found,
"wd"/"wk" onto PG rows.
Answers are cached in years_cache.json.
"""
import argparse
import json
import re
import sys
import time
import unicodedata
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

UA = "NaturalHistoriesCatalogue/1.0 (https://thomashaider.net/natural-histories/)"
WDQS = "https://query.wikidata.org/sparql"
OL = "https://openlibrary.org/search.json"
MIRROR = "https://gutenberg.pglaf.org/cache/epub/{id}/pg{id}.txt"
PLACES = re.compile(r"\b(london|new york|boston|philadelphia|edinburgh|glasgow|dublin|oxford|cambridge|chicago|"
                    r"cincinnati|san francisco|toronto|montreal|sydney|melbourne|calcutta|paris|leipzig|berlin|"
                    r"wien|stuttgart|münchen|munchen|hamburg|dresden|frankfurt|jena|braunschweig|breslau|"
                    r"printed|published|publishers?|verlag|press|& co|and co|sons|brothers|company)\b", re.I)
STOP_AT = re.compile(r"^\s*(contents|table of contents|preface|introduction|chapter|inhalt|vorrede|vorwort|einleitung)\b", re.I)
ROMAN_NUM = re.compile(r"\bM[DCLXVI]{2,}\b")
STOP = set("a an the of and or in on to from with by for at its his her their being into upon "
           "der die das des dem den ein eine einer eines und oder in im am an auf aus von vom zu zur zum mit nach uber fur bei".split())
VOLUME = re.compile(r"(\s*[—–-]+\s*|,\s*|\s+|\.\s*)\(?\b(vol(ume)?|v|part|book|band|bd|teil|theil|th|tome|abt(eilung)?)\.?\s*"
                    r"([0-9]+|[ivxlc]+)\b\)?.*$|\s*\((of \d+|\d+ (vols?|volumes))\)", re.I)
ENTITY = "http://www.wikidata.org/entity/"


def norm(s):
    s = unicodedata.normalize("NFKD", s or "")
    return "".join(c for c in s if not unicodedata.combining(c)).lower()


def words(s):
    return [w for w in re.findall(r"[a-z0-9]+", norm(s)) if w not in STOP]


def base_title(t):
    """Title without subtitle and volume markers: the part worth searching for."""
    t = VOLUME.sub("", t)
    t = re.split(r"\s*[:;\n]|\s+[—–]\s+|(?<=[a-z]{4})\.\s|,\s+(?:or|oder)\b|\s+\((?!of)"
                 r"|\s+(?:Being|Containing|Comprising|Including|Nebst|Enthaltend)\b", t)[0]
    return " ".join(t.split()[:10]).strip(" .,")


def similar(a, b):
    wa, wb = words(a), words(b)
    if not wa or not wb:
        return 0.0
    if wa == wb[:len(wa)] or wb == wa[:len(wb)]:
        return 1.0
    sa, sb = set(wa), set(wb)
    return len(sa & sb) / len(sa | sb)


def life(r):
    """(birth, death) of the first author, from 'Name (1809–1882)'."""
    m = re.search(r"\((\d{3,4})\??–(\d{3,4})?\??\)", (r.get("a") or "").split("; ")[0])
    if not m:
        return r.get("b"), None
    return int(m.group(1)), int(m.group(2)) if m.group(2) else None


def plausible(y, r):
    if y is None or not 1450 <= y <= 1970:
        return False
    born, died = life(r)
    if born and not born + 12 <= y <= (died + 60 if died else born + 80):
        return False
    return True


def surname(r):
    a = (r.get("a") or "").split("; ")[0]
    return re.split(r"[,(]", a)[0].strip()


def get(url, data=None, accept="application/json", timeout=60):
    req = urllib.request.Request(url, data=data, headers={"Accept": accept, "User-Agent": UA})
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return json.load(r)
        except urllib.error.HTTPError as e:
            if e.code not in (429, 500, 502, 503, 504) or attempt == 4:
                raise
            wait = int(e.headers.get("Retry-After") or 0) or 30 * (attempt + 1)
            print(f"  HTTP {e.code}, waiting {wait}s", file=sys.stderr)
            time.sleep(min(wait, 300))


def wikidata(ids, cache, save):
    todo = [i for i in ids if i not in cache]
    for n in range(0, len(todo), 300):
        batch = todo[n:n + 300]
        q = ("SELECT ?pg ?item ?date ?work ?wdate WHERE { VALUES ?pg { %s } ?item wdt:P2034 ?pg . "
             "OPTIONAL { ?item wdt:P577 ?date } OPTIONAL { ?item wdt:P629 ?work . OPTIONAL { ?work wdt:P577 ?wdate } } }"
             % " ".join(f'"{i}"' for i in batch))
        res = get(WDQS, urllib.parse.urlencode({"query": q}).encode(), "application/sparql-results+json", 180)
        for i in batch:
            cache[i] = []
        for b in res["results"]["bindings"]:
            cache[b["pg"]["value"]].append({k: v["value"].replace(ENTITY, "") for k, v in b.items() if k != "pg"})
        save()
        print(f"  wikidata {min(n + 300, len(todo))}/{len(todo)}", file=sys.stderr)
        time.sleep(5)


def year_of(s):
    m = re.match(r"(\d{4})", s or "")
    return int(m.group(1)) if m else None


def from_wikidata(r, hits):
    best = None
    for h in hits:
        for key, src in (("wdate", "work"), ("date", "item")):
            y = year_of(h.get(key))
            if plausible(y, r) and (best is None or (src == "work", -y) > (best[2] == "work", -best[0])):
                best = (y, h, src)
    out = {}
    if hits:
        out["wd"] = hits[0]["item"]
        works = sorted({h["work"] for h in hits if h.get("work")})
        if works:
            out["wk"] = works[0]
    if best:
        out["y"] = best[0]
    return out


def ol_key(r):
    return f"{norm(base_title(r['t']))[:120]}|{norm(surname(r))}"


def openlibrary(rows, cache, delay):
    todo = {}
    for r in rows:
        k = ol_key(r)
        if k not in cache:
            todo[k] = r
    print(f"{len(todo)} Open Library searches", file=sys.stderr)

    def one(item):
        k, r = item
        params = {"title": base_title(r["t"])[:200], "fields": "key,title,author_name,first_publish_year,publish_year", "limit": 8}
        if surname(r):
            params["author"] = surname(r)
        try:
            docs = get(OL + "?" + urllib.parse.urlencode(params)).get("docs", [])
        except Exception as e:  # leave uncached; a rerun retries
            print(f"  ! {r['id']}: {e}", file=sys.stderr)
            docs = None
        time.sleep(delay)
        return k, docs

    # Two requests at a time; with ~2 s per answer that stays under Open Library's 1 request/second.
    with ThreadPoolExecutor(2) as pool:
        for n, (k, docs) in enumerate(pool.map(one, todo.items()), 1):
            if docs is not None:
                cache[k] = docs
            if n % 100 == 0:
                yield n, len(todo)
    yield len(todo), len(todo)


def roman(s):
    vals = {"I": 1, "V": 5, "X": 10, "L": 50, "C": 100, "D": 500, "M": 1000}
    n = 0
    for a, b in zip(s, s[1:] + " "):
        n += -vals[a] if vals.get(b, 0) > vals[a] else vals[a]
    return n


def fetch_heads(ids, folder, delay):
    """First 60 KB of each ebook's text from Gutenberg's mirror (cached); two requests at a time."""
    folder.mkdir(parents=True, exist_ok=True)
    todo = [i for i in ids if not (folder / f"{i}.txt").exists()]
    print(f"{len(todo)} ebook title pages to fetch", file=sys.stderr)

    def one(i):
        req = urllib.request.Request(MIRROR.format(id=i), headers={"User-Agent": UA, "Range": "bytes=0-60000"})
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                (folder / f"{i}.txt").write_bytes(resp.read())
        except Exception as e:  # leave unfetched; a rerun retries
            print(f"  ! {i}: {e}", file=sys.stderr)
        time.sleep(delay)

    with ThreadPoolExecutor(2) as pool:
        for n, _ in enumerate(pool.map(one, todo), 1):
            if n % 100 == 0:
                print(f"  title pages {n}/{len(todo)}", file=sys.stderr)


def from_titlepage(r, text):
    """Publication year from the ebook header or its transcribed title page."""
    m = re.search(r"^Original publication:(.*)$", text, re.M)
    if m:
        ys = [int(y) for y in re.findall(r"\b(1[4-9]\d\d)\b", m.group(1))]
        if ys and 1450 <= min(ys) <= 1970:
            return min(ys)
    start = re.search(r"\*\*\* ?START OF (THE|THIS) PROJECT GUTENBERG EBOOK[^\n]*\n", text, re.I)
    lines = text[start.end() if start else 0:].splitlines()
    in_title = set(re.findall(r"\b1[4-9]\d\d\b", r["t"]))
    copyright, other, seen = [], [], 0
    for line in lines:
        if not line.strip():
            continue
        seen += 1
        if seen > 120 or (seen > 3 and STOP_AT.match(line)):
            break
        s = line.strip()
        if len(s) > 70:  # prose, not a title page line
            continue
        ys = [y for y in re.findall(r"\b(1[4-9]\d\d)\b", s) if y not in in_title]
        ys += [str(roman(x)) for x in ROMAN_NUM.findall(s)]
        ys = [int(y) for y in ys if 1450 <= int(y) <= 1970]
        if not ys:
            continue
        if re.search(r"copyright|entered according to act|copyrighted", s, re.I):
            copyright += ys
        elif PLACES.search(s) or re.fullmatch(r"[\[(]?(1[4-9]\d\d|M[DCLXVI]+)[.\])]?", s):
            other += ys
    for y in sorted(copyright) + other:
        if plausible(y, r):
            return y
    return None


def from_openlibrary(r, docs):
    title = base_title(r["t"])
    sn = norm(surname(r))
    best = None
    for d in docs:
        if sn and not any(sn in norm(a) for a in d.get("author_name", [])):
            continue
        if similar(title, d.get("title", "")) < 0.7:
            continue
        years = sorted(set(d.get("publish_year") or []) | ({d["first_publish_year"]} if d.get("first_publish_year") else set()))
        y = next((y for y in years if plausible(y, r)), None)
        if y and (best is None or y < best):
            best = y
    return {"y": best} if best else {}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--cache", default="years_cache.json")
    ap.add_argument("--delay", type=float, default=1.0, help="seconds between Open Library requests")
    ap.add_argument("--cached-only", action="store_true")
    ap.add_argument("--heads", default="cache/pgheads", help="folder for the ebooks' first 60 KB")
    args = ap.parse_args()

    rows = json.load(open(args.data, encoding="utf-8"))
    cache = json.loads(Path(args.cache).read_text(encoding="utf-8")) if Path(args.cache).exists() else {}
    wd, ol = cache.setdefault("wikidata", {}), cache.setdefault("openlibrary", {})
    save = lambda: Path(args.cache).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
    pg = [r for r in rows if r["src"] == "PG"]

    if not args.cached_only:
        wikidata([r["id"] for r in pg], wd, save)
    found = {}
    for r in pg:
        found[r["id"]] = from_wikidata(r, wd.get(r["id"], []))
    need = [r for r in pg if "y" not in found[r["id"]]]
    if not args.cached_only:
        for n, total in openlibrary(need, ol, args.delay):
            save()
            print(f"  openlibrary {n}/{total}", file=sys.stderr)

    def dated_by_search(r):
        return "y" in found[r["id"]] or from_openlibrary(r, ol.get(ol_key(r), [])).get("y")
    still = [r["id"] for r in pg if not dated_by_search(r)]
    if not args.cached_only:
        fetch_heads(still, Path(args.heads), args.delay)

    counts = {"wikidata": 0, "openlibrary": 0, "titlepage": 0, "lifespan": 0}
    for r in pg:
        for k in ("y", "ys", "wd", "wk"):
            r.pop(k, None)
        f = found[r["id"]]
        r.update({k: v for k, v in f.items() if k in ("wd", "wk")})
        if "y" in f:
            r["y"], r["ys"] = f["y"], "wikidata"
        else:
            y = from_openlibrary(r, ol.get(ol_key(r), [])).get("y")
            if y:
                r["y"], r["ys"] = y, "openlibrary"
        head = Path(args.heads) / f"{r['id']}.txt"
        if not r.get("ys") and head.exists():
            y = from_titlepage(r, head.read_text(encoding="utf-8", errors="replace"))
            if y:
                r["y"], r["ys"] = y, "titlepage"
        born, died = life(r)
        if not r.get("ys") and born and died:
            r["y"], r["ys"] = (born + died) // 2, "lifespan"
        if r.get("ys"):
            counts[r["ys"]] += 1
        else:
            r["y"] = None
    json.dump(rows, open(args.data, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    print(f"Years for {sum(counts.values())}/{len(pg)} Gutenberg works: {counts}", file=sys.stderr)


if __name__ == "__main__":
    main()