#!/usr/bin/env python3
"""Build a corpus of pre-2000 instructional/textbook passages whose full text is
actually downloadable from archive.org.

Two things this had to work around:
  * Lending-library books (collection:inlibrary / printdisabled) return 401 on
    _djvu.txt, which rules out the whole in-copyright 1990s textbook shelf.
  * djvu OCR has no paragraph structure -- every few lines is its own blank-line
    block -- so splitting on blank lines yields nothing. Normalise the whole
    document to one stream, then sentence-pack a sliding window.
"""
import json, os, re, sys, time, urllib.parse, urllib.request

UA = {"User-Agent": "agentatwork/1.0 (+https://agentatwork.xyz)"}
OUT = "corpus.jsonl"

QUERIES = [
    # programmed-instruction / self-study textbooks: uniform declarative prose
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:("individual learning")',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:(textbook)',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:("introduction to")',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:(fundamentals OR principles)',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:("study guide" OR workbook)',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:(handbook OR manual OR primer)',
    'collection:(folkscanomy OR opensource) AND year:[1988 TO 2000] AND title:("course in" OR coursebook)',
    # ESL / controlled-vocabulary teaching text: deliberately flat and uniform
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("english as a second language")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("english for")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:(textbook OR "text book")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("instructional materials")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("teaching guide" OR "teacher guide")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("student manual" OR "learner guide")',
    'collection:(ericarchive) AND year:[1990 TO 1999] AND title:("training manual" OR "training course")',
    'collection:(usgovernmentmirrors) AND year:[1990 TO 1999] AND title:(handbook OR textbook OR "training")',
]

WORD = re.compile(r"[A-Za-z']+")
SENTS = re.compile(r"(?<=[.!?])\s+(?=[A-Z\"'(])")


def get(url, timeout=90):
    for _ in range(3):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=UA), timeout=timeout) as r:
                return r.read()
        except Exception:
            time.sleep(2)
    return None


def search(q, rows=80):
    url = "https://archive.org/advancedsearch.php?" + urllib.parse.urlencode(
        [("q", q), ("fl[]", "identifier"), ("fl[]", "title"), ("fl[]", "year"),
         ("rows", rows), ("output", "json")])
    b = get(url, 60)
    try:
        return json.loads(b)["response"]["docs"] if b else []
    except Exception:
        return []


def djvu_name(ident):
    """The OCR file is named after the item's original title, not its identifier, so
    <identifier>_djvu.txt 404s on anything uploaded with a filename that differs."""
    b = get(f"https://archive.org/metadata/{ident}", 60)
    if not b:
        return None
    try:
        files = json.loads(b).get("files", [])
    except Exception:
        return None
    for f in files:
        if f.get("name", "").endswith("_djvu.txt"):
            return f["name"]
    return None


def normalise(t):
    t = t.replace("\r", "")
    t = re.sub(r"(\w)-\n(\w)", r"\1\2", t)     # de-hyphenate across the line break
    t = re.sub(r"\s+", " ", t)                 # OCR has no paragraphs; make one stream
    return t.strip()


def prose_sentence(s):
    """Keep sentences that look like edited prose, not scan furniture."""
    w = WORD.findall(s)
    if not (4 <= len(w) <= 60):
        return False
    if sum(1 for x in w if len(x) == 1 and x.lower() not in "ai") / len(w) > 0.08:
        return False
    if sum(1 for x in w if x.isupper() and len(x) > 2) / len(w) > 0.15:
        return False
    if len(re.findall(r"\d", s)) / max(1, len(s)) > 0.06:
        return False
    if len(re.findall(r"[^A-Za-z0-9\s.,;:'\"()\-?!%$/&]", s)) / max(1, len(s)) > 0.02:
        return False
    if not re.search(r"\b(the|a|an|is|are|was|were|of|to|and|in|that|this|it|you|we)\b", s, re.I):
        return False
    return True


def passages(text, target=300, lo=240, hi=430):
    """Pack runs of consecutive prose sentences. A run breaks whenever a
    non-prose sentence appears, so passages never straddle a table or a heading."""
    out, buf, n = [], [], 0
    for s in SENTS.split(text):
        s = s.strip()
        if not prose_sentence(s):
            buf, n = [], 0
            continue
        buf.append(s)
        n += len(WORD.findall(s))
        if n >= target:
            p = " ".join(buf)
            if lo <= n <= hi:
                out.append(p)
            buf, n = [], 0
    return out


def main():
    seen = set()
    if os.path.exists(OUT):
        for line in open(OUT):
            try:
                seen.add(json.loads(line)["id"])
            except Exception:
                pass

    uniq = {}
    for q in QUERIES:
        found = search(q)
        print(f"{len(found):4d}  {q[:72]}", file=sys.stderr)
        for d in found:
            uniq.setdefault(d["identifier"], d)
    print(f"unique items: {len(uniq)}", file=sys.stderr)

    kept = 0
    with open(OUT, "a") as fh:
        for i, (ident, d) in enumerate(uniq.items()):
            if ident in seen:
                continue
            name = djvu_name(ident)
            raw = get(f"https://archive.org/download/{ident}/{urllib.parse.quote(name)}") if name else None
            if not raw:
                print(f"[{i+1}/{len(uniq)}] {ident} NO TEXT", file=sys.stderr)
                continue
            ps = passages(normalise(raw.decode("utf-8", "replace")))
            for j, p in enumerate(ps[:10]):     # cap per book so one book cannot flood
                fh.write(json.dumps({"id": ident, "title": d.get("title"),
                                     "year": d.get("year"), "seq": j, "text": p}) + "\n")
                kept += 1
            fh.flush()
            print(f"[{i+1}/{len(uniq)}] {ident} {d.get('year')} -> {len(ps)} passages", file=sys.stderr)
    print(f"TOTAL kept: {kept}", file=sys.stderr)


if __name__ == "__main__":
    main()