#!/usr/bin/env python3
"""
Sort the catalogue's keywords into kinds for the keyword overview page.

Usage:
    python keyword_types.py        # after keywords.py (and group.py); -> overview/keywords.json

An LLM on innkube (qwen38-27b, Uni Passau) assigns each distinct keyword one kind: place,
animals, plants, landscape, peoples, activity, science, period, form, or other. Answers are
cached in keyword_types_cache.json. Needs INNKUBE_API_KEY.
"""
import argparse
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 = 150
KINDS = [  # (key, label, description given to the LLM)
    ("place", "Places and regions", "a named country, region, city, island, river, mountain range, sea or continent"),
    ("animals", "Animals", "animals or groups of animals (birds, insects, whales, game, fauna) and their study (ornithology, entomology)"),
    ("plants", "Plants and fungi", "plants, trees, flowers, fungi, flora, and their study (botany)"),
    ("landscape", "Landscapes and nature", "kinds of landscape, environment, weather or natural phenomena (mountains, forests, deserts, ice, seasons, wilderness, scenery)"),
    ("peoples", "Peoples and society", "peoples, cultures, customs, religion, social and economic life, settlements, architecture"),
    ("activity", "Travel and activities", "journeys, voyages, expeditions, exploration, hunting, fishing, mountaineering, farming, collecting and other activities"),
    ("science", "Sciences and ideas", "sciences and ideas other than botany and zoology (geology, evolution, natural history, ecology, geography, medicine, philosophy)"),
    ("period", "Periods and events", "historical periods, centuries, wars and events"),
    ("form", "Forms and genres", "the kind of text (letters, diaries, essays, sketches, guidebook, juvenile literature, illustrations, poetry, biography)"),
    ("other", "Other", "anything else"),
]
PROMPT = ("Sort each numbered keyword from a catalogue of historical natural histories and travel writing into exactly one kind:\n"
          + "\n".join(f"- {k}: {d}" for k, _, d in KINDS)
          + "\nAnswer only with JSON mapping each number to one kind key.\n\n")


def classify(words, cache_path, workers):
    cache = json.loads(Path(cache_path).read_text(encoding="utf-8")) if Path(cache_path).exists() else {}
    todo = [w for w in words if w not in cache]
    batches = [todo[i:i + BATCH] for i in range(0, len(todo), BATCH)]
    print(f"{len(todo)} keywords to sort ({len(batches)} requests)", file=sys.stderr)
    valid = {k for k, _, _ in KINDS}

    def ask(batch):
        prompt = PROMPT + "\n".join(f"{n}. {w}" for n, w in enumerate(batch, 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"})
        try:
            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)
            return batch, json.loads(content[content.index("{"):content.rindex("}") + 1])
        except Exception as e:  # leave uncached; a rerun retries
            print(f"  ! batch failed: {e}", file=sys.stderr)
            return batch, None

    with ThreadPoolExecutor(workers) as pool:
        for n, (batch, answer) in enumerate(pool.map(ask, batches), 1):
            if answer:
                for i, w in enumerate(batch, 1):
                    if answer.get(str(i)) in valid:
                        cache[w] = answer[str(i)]
            Path(cache_path).write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
            print(f"  {n}/{len(batches)}", file=sys.stderr)
    return cache


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--cache", default="keyword_types_cache.json")
    ap.add_argument("-o", "--out", default="overview/keywords.json")
    ap.add_argument("--workers", type=int, default=4)
    args = ap.parse_args()
    rows = json.load(open(args.data, encoding="utf-8"))
    words = sorted({k.lower() for r in rows for k in r.get("kw", [])})
    cache = classify(words, args.cache, args.workers)
    out = {"kinds": [{"key": k, "label": l} for k, l, _ in KINDS], "kw": {w: cache.get(w, "other") for w in words}}
    json.dump(out, open(args.out, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    print(f"Wrote {len(words)} keywords to {args.out} ({sum(1 for w in words if w not in cache)} unsorted)", file=sys.stderr)


if __name__ == "__main__":
    main()