"""Generate full-length MODERN AI academic text (Gemini) to fix the detector's
blocker: no freely-available modern AI academic corpus.

Seeds come from the HUMAN academic texts (data/calibration/mage_sci.jsonl) so
every AI doc is TOPIC-MATCHED to a real human doc — the detector then learns
authorship, not topic. Three realistic modes mirror how students use AI:
  write   : "write this academic section about <topic>"   (pure generation)
  rewrite : "rewrite this human passage academically"      (AI-assisted/humanised)
  continue: "continue this academic passage"               (continuation)

Free-tier safe: gemini-flash-latest, exponential backoff on 429, incremental
+ resumable writes. Key from $GEMINI_API_KEY or data/.gemini_key (gitignored).

Output: data/calibration/gen_ai.jsonl  ({text, y:0, src_model, mode, seed_id})
Run: python scripts/gen_ai_corpus.py [n_per_mode]   (default 50)
"""
import json, os, re, sys, time, urllib.request, urllib.error

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
HUMAN = os.path.join(ROOT, "data", "calibration", "mage_sci.jsonl")
OUT = os.path.join(ROOT, "data", "calibration", "gen_ai.jsonl")
MODEL = os.environ.get("GEMINI_MODEL", "gemini-flash-latest")
N_PER_MODE = int(sys.argv[1]) if len(sys.argv) > 1 else 50


def get_key():
    k = os.environ.get("GEMINI_API_KEY")
    if not k:
        p = os.path.join(ROOT, "data", ".gemini_key")
        if os.path.exists(p):
            k = open(p).read().strip()
    if not k:
        sys.exit("no GEMINI_API_KEY (env or data/.gemini_key)")
    return k


KEY = get_key()
URL = (f"https://generativelanguage.googleapis.com/v1beta/models/"
       f"{MODEL}:generateContent?key={KEY}")


def gemini(prompt, temp=0.9, max_tokens=900, tries=6):
    body = json.dumps({
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": temp, "maxOutputTokens": max_tokens},
    }).encode()
    delay = 30.0
    for t in range(tries):
        try:
            req = urllib.request.Request(URL, data=body,
                                         headers={"Content-Type": "application/json"})
            r = json.load(urllib.request.urlopen(req, timeout=60))
            cand = r.get("candidates", [{}])[0]
            parts = cand.get("content", {}).get("parts", [{}])
            return "".join(p.get("text", "") for p in parts).strip()
        except urllib.error.HTTPError as e:
            msg = e.read().decode() if e.code == 429 else ""
            if e.code == 429:                      # free-tier per-minute limit
                # Gemini puts the retry time in the body: "Please retry in 28.5s"
                m = re.search(r"retry in ([\d.]+)s", msg)
                ra = e.headers.get("Retry-After")
                wait = float(m.group(1)) + 2 if m else (float(ra) if ra else delay)
                print(f"    429 — waiting {wait:.0f}s for quota window", flush=True)
                time.sleep(wait)
            else:
                print(f"    HTTP {e.code}: {e.read().decode()[:120]}", flush=True)
                time.sleep(delay); delay = min(delay * 1.5, 90)
        except Exception as ex:
            print(f"    ERR {str(ex)[:100]}", flush=True)
            time.sleep(delay); delay = min(delay * 1.5, 90)
    return ""


def topic_of(text):
    """A short topic handle from a human academic doc (its opening)."""
    return " ".join(text.split()[:40])


PROMPTS = {
    "write": ("Write a detailed, formal section of an academic research paper "
              "(~400 words) on the topic introduced by the following excerpt. "
              "Use rigorous academic style, no headings, prose only:\n\n{seed}"),
    "rewrite": ("Rewrite the following academic passage entirely in your own "
                "words, preserving the meaning and academic tone (~350 words), "
                "prose only:\n\n{seed}"),
    "continue": ("Continue the following academic passage for about 350 more "
                 "words in the same formal academic style, prose only:\n\n{seed}"),
}


def done_ids():
    if not os.path.exists(OUT):
        return set()
    return {(json.loads(l)["mode"], json.loads(l)["seed_id"])
            for l in open(OUT, encoding="utf-8")}


def main():
    humans = [json.loads(l) for l in open(HUMAN, encoding="utf-8")
              if json.loads(l)["y"] == 1]
    seeds = [(i, h["text"]) for i, h in enumerate(humans) if len(h["text"]) > 600]
    already = done_ids()
    print(f"{len(seeds)} human seeds, model={MODEL}, "
          f"{N_PER_MODE}/mode, resume (have {len(already)})")
    out = open(OUT, "a", encoding="utf-8")
    made = 0
    for mode, tmpl in PROMPTS.items():
        n = 0
        for sid, seed in seeds:
            if n >= N_PER_MODE:
                break
            if (mode, sid) in already:
                n += 1
                continue
            seed_txt = topic_of(seed) if mode == "write" else seed[:1600]
            txt = gemini(tmpl.format(seed=seed_txt))
            if len(txt) < 500:                     # too short / blocked / empty
                continue
            out.write(json.dumps({"text": txt, "y": 0, "src_model": "ai_gemini",
                                  "mode": mode, "seed_id": sid}) + "\n")
            out.flush()
            n += 1; made += 1
            print(f"  [{mode}] {n}/{N_PER_MODE}  ({len(txt)} chars)", flush=True)
            time.sleep(4.5)                        # free-tier pacing
    out.close()
    total = sum(1 for _ in open(OUT, encoding="utf-8"))
    print(f"\ngenerated {made} this run; {total} total -> {OUT}")


if __name__ == "__main__":
    main()