#!/usr/bin/env python3
"""
Pick the natural histories and travel writing out of EcoCor (https://ecocor.org).

Usage:
    python ecocor_nature_travel.py                  # -> ecocor_nature_travel.csv

EcoCor's German and English corpora are mostly fiction, with plant and animal
annotations. An LLM on innkube (qwen38-27b, Uni Passau) labels each text natural_history,
travel or neither from its author, title and year; labels are cached in
ecocor_labels.json. Needs INNKUBE_API_KEY.
"""
import argparse
import csv
import json
import os
import re
import sys
import urllib.request
from pathlib import Path

API = "https://ecocor.org/api"
INNKUBE = "https://llms.innkube.fim.uni-passau.de/v1/chat/completions"
LLM_MODEL = "qwen38-27b"
# Texts the LLM labelled travel but which are fiction (Dauthendey's Asian novellas/stories).
EXCLUDE = {"1909_Dauthendey_Der-Knabe-auf", "1909_Dauthendey_Die-Abendglocken-vom",
           "1911_Dauthendey_Die-acht-Gesichter"}
PROMPT = (
    "Label each numbered work (author: title, year). Use 'natural_history' for non-fiction about "
    "animals, plants, landscapes or nature (including nature sketches and nature essays), 'travel' "
    "for non-fiction accounts of journeys or of places the author visited, both labels if both fit, "
    "and 'neither' for fiction, poetry, drama and everything else (animal stories and novels set "
    "abroad are fiction). Answer only with JSON mapping each number to a list of labels.\n\n")


def get(url):
    with urllib.request.urlopen(urllib.request.Request(url, headers={"Accept": "application/json"}),
                                timeout=120) as r:
        return json.load(r)


def label(texts, cache_path):
    cache = json.loads(Path(cache_path).read_text(encoding="utf-8")) if Path(cache_path).exists() else {}
    todo = [t for t in texts if t["name"] not in cache]
    for i in range(0, len(todo), 40):
        batch = todo[i:i + 40]
        lines = [f"{n}. {', '.join(a['name'] for a in t['authors'])}: {t['title']}, {t['referenceYear']}"
                 for n, t in enumerate(batch, 1)]
        body = {"model": LLM_MODEL, "messages": [{"role": "user", "content": PROMPT + "\n".join(lines)}],
                "temperature": 0, "max_tokens": 3000, "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])
        for n, t in enumerate(batch, 1):
            cache[t["name"]] = [l for l in answer.get(str(n), []) if l in ("natural_history", "travel")]
        Path(cache_path).write_text(json.dumps(cache, ensure_ascii=False, indent=1), encoding="utf-8")
    return cache


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-o", "--out", default="ecocor_nature_travel.csv")
    ap.add_argument("--labels", default="ecocor_labels.json")
    args = ap.parse_args()

    texts = [dict(t, corpus=c) for c in ("de", "en") for t in get(f"{API}/corpora/{c}/texts")]
    labels = label(texts, args.labels)
    rows = []
    for t in texts:
        cats = labels.get(t["name"], [])
        if not cats or t["name"] in EXCLUDE:
            continue
        m = t.get("metrics", {})
        rows.append({
            "name": t["name"], "corpus": t["corpus"], "title": t["title"],
            "author": "; ".join(a["name"] for a in t["authors"]),
            "year": t.get("referenceYear") or "", "category": "+".join(sorted(cats)),
            "animals": m.get("numOfAnimals", ""), "plants": m.get("numOfPlants", ""),
            "words": m.get("numOfWords", ""),
            "url": f"https://ecocor.org/corpora/{t['corpus']}/{t['name']}",
        })
    with open(args.out, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(sorted(rows, key=lambda r: r["year"]))
    print(f"{len(rows)} of {len(texts)} EcoCor texts kept -> {args.out}", file=sys.stderr)
    for r in rows:
        print(f"  {r['corpus']} {r['year']} {r['category']:24s} {r['author'][:25]} | {r['title'][:60]}",
              file=sys.stderr)


if __name__ == "__main__":
    main()