"""Build the study population and the control set from the 1f916.ai corpus.

Applies rule.py exactly. Writes population.csv, control.csv and
extension-urlfield.csv, and prints the counts. No network.

  python research/agent-site-mortality/extract.py
"""
import csv
import json
import os
import re
import sys
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import rule  # noqa: E402

RAW = os.path.expanduser("~/data/1f916")
OUT = os.path.dirname(os.path.abspath(__file__))

# A URL in prose ends at whitespace or at punctuation that is plainly the
# sentence's and not the URL's. Trailing `)` is stripped only when the URL
# carries no `(`, so a wikipedia-style parenthesised path survives.
URL_RE = re.compile(r"https?://[^\s<>\"'`\\]+", re.I)
TRAILING = ".,;:!?'\"”’>*_)]}"


def clean(u):
    while u and u[-1] in TRAILING:
        if u[-1] == ")" and "(" in u:
            break
        u = u[:-1]
    return u


def host_of(u):
    m = re.match(r"https?://([^/?#]+)", u, re.I)
    if not m:
        return None
    h = m.group(1).lower()
    h = h.split("@")[-1]          # strip any userinfo
    h = h.split(":")[0]           # strip the port
    if h.startswith("www."):
        h = h[4:]
    if not h or h == "localhost":
        return None
    if re.fullmatch(r"[\d.]+", h) or ":" in h:   # IPv4 / IPv6 literal
        return None
    if "." not in h:
        return None
    return h


def rows(name):
    with open(os.path.join(RAW, name + ".jsonl")) as fh:
        for line in fh:
            if line.strip():
                yield json.loads(line)


def main():
    # host -> {citizen -> link count}; and the corpus-wide last activity per citizen
    by_host = defaultdict(lambda: defaultdict(int))
    urlfield = defaultdict(lambda: defaultdict(int))
    last_seen = {}
    first_link = {}
    last_link = {}
    n_posts = n_comments = 0

    for kind in ("posts", "comments"):
        for r in rows(kind):
            if kind == "posts":
                n_posts += 1
            else:
                n_comments += 1
            author = r.get("author")
            ts = r.get("created_at")
            if author and ts:
                if author not in last_seen or ts > last_seen[author]:
                    last_seen[author] = ts
            if not author:
                continue
            for raw_url in URL_RE.findall(r.get("body") or ""):
                h = host_of(clean(raw_url))
                if not h:
                    continue
                by_host[h][author] += 1
                if ts:
                    if h not in first_link or ts < first_link[h]:
                        first_link[h] = ts
                    if h not in last_link or ts > last_link[h]:
                        last_link[h] = ts
            # the `url` field of a link post: NOT in the sealed rule, kept
            # separately as a declared extension
            if kind == "posts" and r.get("url"):
                h = host_of(clean(r["url"]))
                if h:
                    urlfield[h][author] += 1

    population, control = [], []
    for h, authors in by_host.items():
        if rule.is_shared(h) or h == rule.SELF_HOST:
            continue
        n_cit = len(authors)
        n_links = sum(authors.values())
        if n_cit == rule.MAX_CITIZENS_PERSONAL and n_links >= rule.MIN_LINKS_PERSONAL:
            owner = next(iter(authors))
            population.append({
                "host": h, "owner": owner, "links": n_links,
                "first_link": first_link.get(h, ""), "last_link": last_link.get(h, ""),
                "owner_last_seen": last_seen.get(owner, ""),
            })
        elif n_cit >= rule.MIN_CITIZENS_CONTROL:
            control.append({
                "host": h, "citizens": n_cit, "links": n_links,
                "first_link": first_link.get(h, ""), "last_link": last_link.get(h, ""),
            })

    ext = []
    for h, authors in urlfield.items():
        if rule.is_shared(h) or h == rule.SELF_HOST:
            continue
        if len(authors) == 1 and sum(authors.values()) >= rule.MIN_LINKS_PERSONAL:
            owner = next(iter(authors))
            ext.append({"host": h, "owner": owner, "links": sum(authors.values()),
                        "already_in_population": any(p["host"] == h for p in population)})

    population.sort(key=lambda r: (-r["links"], r["host"]))
    control.sort(key=lambda r: (-r["citizens"], r["host"]))
    ext.sort(key=lambda r: (-r["links"], r["host"]))

    for name, data in (("population", population), ("control", control),
                       ("extension-urlfield", ext)):
        path = os.path.join(OUT, name + ".csv")
        with open(path, "w", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=list(data[0].keys()) if data else ["host"])
            w.writeheader()
            w.writerows(data)
        print(f"{name}: {len(data)} hosts -> {os.path.basename(path)}")

    print(f"corpus: {n_posts} posts, {n_comments} comments, "
          f"{len(by_host)} distinct hosts in body text, "
          f"{len(last_seen)} citizens who wrote anything")


if __name__ == "__main__":
    main()