#!/usr/bin/env python3
"""
Add 3–6 English keywords per work to overview/data.json, so a reader can see at a glance
what each work is about.

Usage:
    python keywords.py            # after build_overview.py and geolocate.py

An LLM on innkube (qwen38-27b, Uni Passau) reads each work's author, title, year and
whatever description exists (LoC subjects, DTA genre, Wikisource source and summary) and
returns keywords. Answers are cached in keywords_cache.json, so reruns only ask for new
works. Needs INNKUBE_API_KEY.
"""
import argparse
import hashlib
import json
import os
import re
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

INNKUBE = "https://llms.innkube.fim.uni-passau.de/v1/chat/completions"
LLM_MODEL = "qwen38-27b"
BATCH = 30
PROMPT = (
    "Each numbered line describes a historical work of travel writing or natural history. For each, "
    "give 3 to 6 short English keywords that tell a reader what the work is about: its main subjects "
    "(animals, plants, landscapes, peoples, activities such as whaling or mountaineering, events) and "
    "its main region. Do not use the words travel, natural history, book, journey or nature on their "
    "own. Use lowercase except for proper nouns. Answer only with JSON mapping each number to a list "
    "of keywords.\n\n")


def describe(r):
    parts = [r.get("a") or "anonymous", r["t"]]
    if r.get("y") or r.get("b"):
        parts.append(f"published {r['y']}" if r.get("y") else f"author born {r['b']}")
    if r.get("sub"):
        parts.append("LoC subjects: " + r["sub"][:300])
    elif r.get("sum"):
        parts.append("description: " + r["sum"][:200])
    if r["src"] in ("DTA", "WS") and r.get("g"):
        parts.append(("DTA genre: " if r["src"] == "DTA" else "appeared in: ") + r["g"][:80])
    return " | ".join(parts)


def ask(lines):
    prompt = PROMPT + "\n".join(f"{n}. {l}" for n, l in enumerate(lines, 1))
    body = {"model": LLM_MODEL, "messages": [{"role": "user", "content": prompt}], "temperature": 0,
            "max_tokens": 4000, "chat_template_kwargs": {"enable_thinking": False}}
    req = urllib.request.Request(INNKUBE, data=json.dumps(body).encode(), headers={
        "Authorization": "Bearer " + os.environ["INNKUBE_API_KEY"], "Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        content = json.load(r)["choices"][0]["message"]["content"]
    content = re.sub(r"<think>.*?</think>|```(json)?", "", content, flags=re.S)
    answer = json.loads(content[content.index("{"):content.rindex("}") + 1])
    return [[k.strip() for k in answer.get(str(n), []) if isinstance(k, str) and k.strip()][:6]
            for n in range(1, len(lines) + 1)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--cache", default="keywords_cache.json")
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--cached-only", action="store_true", help="use only keywords already in the cache")
    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 {}
    key = lambda d: hashlib.sha1(d.encode()).hexdigest()[:16]
    descs = {id(r): describe(r) for r in rows}
    todo = list(dict.fromkeys(d for d in descs.values() if key(d) not in cache))
    if args.cached_only:
        todo = []
    batches = [todo[i:i + BATCH] for i in range(0, len(todo), BATCH)]
    print(f"{len(todo)} works need keywords ({len(batches)} requests)", file=sys.stderr)

    def run(batch):
        try:
            return batch, ask(batch)
        except Exception as e:  # leave uncached; a rerun retries
            print(f"  ! batch failed: {e}", file=sys.stderr)
            return batch, None

    with ThreadPoolExecutor(args.workers) as pool:
        for n, (batch, result) in enumerate(pool.map(run, batches), 1):
            if result:
                for d, kws in zip(batch, result):
                    cache[key(d)] = kws
            if n % 10 == 0 or n == len(batches):
                Path(args.cache).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
                print(f"  {n}/{len(batches)}", file=sys.stderr)

    missing = 0
    for r in rows:
        kws = cache.get(key(descs[id(r)]))
        if kws:
            r["kw"] = kws
        else:
            missing += 1
    json.dump(rows, open(args.data, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    print(f"Keywords for {len(rows) - missing}/{len(rows)} works", file=sys.stderr)


if __name__ == "__main__":
    main()