#!/usr/bin/env python3
"""Measure AI-answer citation share for a domain, across four answer engines.

    python3 citation_meter.py --config tines.json
    python3 citation_meter.py --config tines.json --dry-run    # shape only, no keys needed

Why this exists
---------------
Semrush reports whether an AI Overview *appeared* on a keyword
(serp_ai_overview_keywords). It does not report whether your domain was *named
inside the answer*. Those are different questions and only the second one tells
you whether you are winning the surface.

The paid tools that do answer it (Profound, Peec, Otterly, Scrunch, Semrush's
own AI toolkit) all sit behind subscriptions. But every engine worth measuring
already exposes its own citations through a documented public API, so the meter
can be built directly and run for cents.

Engines, and what each returns
------------------------------
  perplexity  api.perplexity.ai/chat/completions          citations / search_results
  openai      api.openai.com/v1/responses + web_search    url_citation annotations
  anthropic   api.anthropic.com/v1/messages + web_search  web_search_tool_result
  gemini      generativelanguage.googleapis.com + search  groundingMetadata

Each engine runs only if its key is present, so this degrades gracefully:
  PERPLEXITY_API_KEY  OPENAI_API_KEY  ANTHROPIC_API_KEY  GEMINI_API_KEY

Or set OPENROUTER_API_KEY alone and all four run through OpenRouter. Native keys
win where both are present, because they return richer citation payloads.

Config file
-----------
    {
      "domain": "tines.com",
      "competitors": ["torq.io", "n8n.io", "swimlane.com", "workato.com"],
      "prompts": [
        "What is the best SOAR platform for a small security team?",
        "Tines vs n8n for security automation",
        "How do I automate phishing triage?"
      ]
    }

Output
------
Per engine: how many prompts cited you, how many cited each competitor, and the
full URL list per prompt so any number here can be traced back to a raw response.

Every run appends to a history directory (default ./runs/), and each run is
compared against the previous one, so the second run onward gives you a trend
rather than a snapshot. Read the series any time without spending a token:

    python3 citation_meter.py --config tines.json --trend

The denominator only means something if the prompt set holds still, so each run
records a hash of its prompts. Change the prompts and the trend view splits the
series at that point rather than quietly comparing different questions.

Weekly, via cron:

    0 9 * * 1  cd /path/to/tools && /usr/bin/python3 citation_meter.py --config tines.json

Arms: holding everything still except one thing
-----------------------------------------------
A config can carry named "arms" instead of a flat prompt list. Each arm is a
different way of asking the same buyer about the same criteria, so a difference
in who gets cited is attributable to the framing and to nothing else.

    "arms": {
      "soar":       {"label": "...", "prompts": [...]},
      "automation": {"label": "...", "prompts": [...]}
    }

    python3 citation_meter.py --config tines-framing-test.json --all-arms
    python3 citation_meter.py --config tines-framing-test.json --arm soar --trend

Each arm keeps its own history series and its own prompt hash, so arms never
contaminate each other's trend. --all-arms adds a comparison table across arms.
An arm that names a competitor in every prompt has primed that competitor's
number; read your own domain in that arm, not theirs.

Response shapes drift between API versions. Rather than couple to a schema, the
parser walks the whole JSON and collects anything that looks like a URL, then
matches on registrable domain. That is deliberately dumb and hard to break.
"""

import argparse
import hashlib
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from collections import defaultdict
from datetime import date, timezone, datetime
from pathlib import Path

URL_RE = re.compile(r'https?://[^\s"\'<>\\)\]}]+', re.I)
TIMEOUT = 90


# --------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------

def post(url: str, payload: dict, headers: dict) -> dict:
    body = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=body, method="POST",
                                 headers={"content-type": "application/json", **headers})
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        return json.loads(r.read().decode())


def harvest_urls(node) -> list[str]:
    """Walk any JSON structure and pull out every URL-looking string.

    Schema-independent on purpose: an engine can rename its citation field and
    this keeps working.
    """
    found = []
    if isinstance(node, str):
        found += URL_RE.findall(node)
    elif isinstance(node, dict):
        for v in node.values():
            found += harvest_urls(v)
    elif isinstance(node, list):
        for v in node:
            found += harvest_urls(v)
    return found


def registrable(url: str) -> str:
    """Host without www. Good enough for citation matching; not a PSL parser."""
    m = re.match(r'https?://([^/:?#]+)', url, re.I)
    if not m:
        return ""
    return m.group(1).lower().removeprefix("www.")


def cited(urls: list[str], domain: str) -> bool:
    d = domain.lower().removeprefix("www.")
    return any(h == d or h.endswith("." + d) for h in (registrable(u) for u in urls))


def prompt_hash(prompts: list[str]) -> str:
    """Identity of a prompt set. A changed set means a changed denominator."""
    return hashlib.sha256("\u0000".join(prompts).encode()).hexdigest()[:12]


def run_path(runs_dir: Path, domain: str, arm: str | None, stamp: str) -> Path:
    slug = domain.replace(".", "-")
    return runs_dir / f"citations-{slug}-{arm + '-' if arm else ''}{stamp}.json"


def load_history(runs_dir: Path, domain: str, arm: str | None = None) -> list[dict]:
    """Every prior run for this domain, oldest first.

    Arms are separate series. Without this, three framings of the same question
    would land in one history and trend_table would compare rows that are not
    comparable, which is the exact failure it exists to warn about.
    """
    slug = domain.replace(".", "-")
    out = []
    for f in sorted(runs_dir.glob(f"citations-{slug}-*.json")):
        try:
            rec = json.loads(f.read_text())
        except (OSError, json.JSONDecodeError):
            print(f"  skipping unreadable run {f.name}", file=sys.stderr)
            continue
        if rec.get("arm") != arm:         # None matches only the single-arm runs
            continue
        out.append(rec)
    return sorted(out, key=lambda r: r.get("run_at", ""))


def trend_table(history: list[dict], domain: str) -> str:
    """Citation share over time, per engine, with the change since last run."""
    if not history:
        return "\nno history yet. this is run one; the trend starts at run two."
    engines = sorted({e for r in history for e in r.get("engines", {})})
    out = [f"\ncitation share over time, {domain}", "=" * 72,
           f"{'run':<22}{'prompts':>8}  " + "  ".join(f"{e:>11}" for e in engines)]
    prev, split_warned = {}, set()
    for r in history:
        stamp = r.get("run_at", "?")[:16].replace("T", " ")
        h = r.get("prompt_set", "?")
        cells, n = [], 0
        for e in engines:
            d = r.get("engines", {}).get(e)
            if not d or d.get("citation_share") is None:
                cells.append(f"{'-':>11}"); continue
            n = max(n, d.get("prompts_run", 0))
            cur = d["citation_share"]
            if e in prev and prev[e] is not None:
                delta = cur - prev[e]
                arrow = "+" if delta > 0 else ""
                cells.append(f"{cur:.0%} ({arrow}{delta:+.0%})".rjust(11).replace("(+-", "(-"))
            else:
                cells.append(f"{cur:.0%}".rjust(11))
            prev[e] = cur
        out.append(f"{stamp:<22}{n:>8}  " + "  ".join(cells))
        if h not in split_warned and len(split_warned) > 0:
            out.append(f"{'':22}{'':8}  prompt set changed here, series restarts")
        split_warned.add(h)
    out.append("=" * 72)
    if len(split_warned) > 1:
        out.append("WARNING: more than one prompt set in this history. Rows either side of a")
        out.append("change are answering different questions and are not comparable.")
    return "\n".join(out)


# --------------------------------------------------------------------------
# engines. each returns (list_of_urls, raw_response) or raises
# --------------------------------------------------------------------------

def ask_perplexity(prompt: str, key: str, model="sonar-pro"):
    r = post("https://api.perplexity.ai/chat/completions",
             {"model": model, "messages": [{"role": "user", "content": prompt}]},
             {"authorization": f"Bearer {key}"})
    return harvest_urls(r), r


def ask_openai(prompt: str, key: str, model="gpt-5.2"):
    r = post("https://api.openai.com/v1/responses",
             {"model": model, "input": prompt, "tools": [{"type": "web_search"}]},
             {"authorization": f"Bearer {key}"})
    return harvest_urls(r), r


def ask_anthropic(prompt: str, key: str, model="claude-sonnet-5"):
    r = post("https://api.anthropic.com/v1/messages",
             {"model": model, "max_tokens": 1500,
              "messages": [{"role": "user", "content": prompt}],
              "tools": [{"type": "web_search_20250305", "name": "web_search",
                         "max_uses": 5}]},
             {"x-api-key": key, "anthropic-version": "2023-06-01"})
    return harvest_urls(r), r


def ask_gemini(prompt: str, key: str, model="gemini-2.5-flash"):
    r = post(f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}",
             {"contents": [{"parts": [{"text": prompt}]}],
              "tools": [{"google_search": {}}]},
             {})
    return harvest_urls(r), r


ENGINES = {
    "perplexity": ("PERPLEXITY_API_KEY", ask_perplexity),
    "openai":     ("OPENAI_API_KEY",     ask_openai),
    "anthropic":  ("ANTHROPIC_API_KEY",  ask_anthropic),
    "gemini":     ("GEMINI_API_KEY",     ask_gemini),
}

# One key instead of four. OpenRouter is OpenAI-compatible and proxies all of
# these, so OPENROUTER_API_KEY alone gets the meter running. The trade is real:
# native APIs return richer citation payloads (Anthropic's search results,
# Gemini's groundingMetadata), whereas OpenRouter returns whatever the upstream
# model puts in the message body plus its own annotations. Use native keys when
# you have them; use this when you want to be measuring in ten minutes.
OPENROUTER_MODELS = {
    "perplexity": "perplexity/sonar-pro",
    "openai":     "openai/gpt-5.2:online",
    "anthropic":  "anthropic/claude-sonnet-5:online",
    "gemini":     "google/gemini-2.5-flash:online",
}


def ask_openrouter(prompt: str, key: str, model: str):
    r = post("https://openrouter.ai/api/v1/chat/completions",
             {"model": model, "messages": [{"role": "user", "content": prompt}],
              "plugins": [{"id": "web"}]},
             {"authorization": f"Bearer {key}",
              "http-referer": "https://www.skilltrade.marketing",
              "x-title": "citation-meter"})
    return harvest_urls(r), r


# --------------------------------------------------------------------------

def run(cfg: dict, dry: bool, prompts: list[str] | None = None,
        arm: str | None = None) -> dict:
    domain = cfg["domain"]
    competitors = cfg.get("competitors", [])
    prompts = prompts if prompts is not None else cfg["prompts"]

    # sparktoro.py writes cfg["audience"], which says what share of THIS buyer
    # actually uses each assistant. When it is present, meter the engines the
    # buyer over-indexes on first, and keep the weights for the summary.
    audience = cfg.get("audience") or {}
    weights = {e: w["audience_share"] for e, w in audience.get("engines", {}).items()}
    order = sorted(ENGINES, key=lambda e: -weights.get(e, 0)) if weights else list(ENGINES)

    router_key = os.environ.get("OPENROUTER_API_KEY")
    active = {}
    for name in order:
        env, fn = ENGINES[name]
        key = os.environ.get(env)
        if key:
            active[name] = (key, fn)
        elif router_key:
            model = OPENROUTER_MODELS[name]
            active[name] = (router_key,
                            lambda p, k, m=model: ask_openrouter(p, k, m))
            print(f"  {name}: via OpenRouter ({model})", file=sys.stderr)
        else:
            print(f"  skipping {name}: neither {env} nor OPENROUTER_API_KEY set", file=sys.stderr)
    if not active and not dry:
        sys.exit("no keys found. set OPENROUTER_API_KEY for all four, or any "
                 "native key, or pass --dry-run")

    results = {"domain": domain, "competitors": competitors, "arm": arm,
               "run_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
               "prompt_set": prompt_hash(prompts), "prompt_count": len(prompts),
               "_prompts": prompts,
               "audience": audience, "engines": {}, "detail": []}

    for name, (key, fn) in (active.items() if not dry else []):
        hits, comp_hits, rows = 0, defaultdict(int), []
        for i, prompt in enumerate(prompts, 1):
            try:
                urls, raw = fn(prompt, key)
            except urllib.error.HTTPError as e:
                print(f"  {name} p{i}: HTTP {e.code} {e.read()[:180]!r}", file=sys.stderr)
                continue
            except Exception as e:                      # noqa: BLE001
                print(f"  {name} p{i}: {type(e).__name__} {e}", file=sys.stderr)
                continue
            urls = sorted(set(urls))
            you = cited(urls, domain)
            hits += you
            for c in competitors:
                comp_hits[c] += cited(urls, c)
            rows.append({"prompt": prompt, "cited": you,
                         "competitors_cited": [c for c in competitors if cited(urls, c)],
                         "urls": urls})
            print(f"  {name} [{i}/{len(prompts)}] {'CITED' if you else '  -  '}  {prompt[:58]}")
            time.sleep(1)                                # be polite
        n = len(rows) or 1
        results["engines"][name] = {
            "prompts_run": len(rows),
            "cited": hits,
            "citation_share": round(hits / n, 3),
            "competitor_share": {c: round(v / n, 3) for c, v in comp_hits.items()},
        }
        results["detail"].append({"engine": name, "rows": rows})

    if dry:
        results["engines"] = {e: {"prompts_run": 0, "cited": 0, "citation_share": None,
                                  "competitor_share": {c: None for c in competitors}}
                              for e in order}
        results["note"] = "dry run: no engines called, shape only"
    return results


def summarise(r: dict) -> str:
    dom, comps = r["domain"], r["competitors"]
    # Brands as rows, engines as columns. The other way round stops fitting a
    # terminal once the competitor set spans two categories.
    engines = list(r["engines"])
    brands = [dom] + comps
    bw = max(len(b) for b in brands) + 1
    head = f"\ncitation share, {r['run_at']}"
    if r.get("arm"):
        head += f"  [arm: {r['arm']}]"
    out = [head, "=" * (bw + 13 * len(engines)),
           f"{'brand':<{bw}}" + "".join(f"{e:>13}" for e in engines)]

    def cell(engine_name: str, brand: str) -> str:
        e = r["engines"][engine_name]
        v = e["citation_share"] if brand == dom else e["competitor_share"].get(brand)
        return "n/a" if v is None else f"{v:.0%}"

    for b in brands:
        mark = "*" if b == dom else " "
        out.append(f"{mark}{b:<{bw - 1}}" + "".join(f"{cell(e, b):>13}" for e in engines))
    ran = {e: r["engines"][e]["prompts_run"] for e in engines}
    out.append("=" * (bw + 13 * len(engines)))
    out.append(f"{'prompts run':<{bw}}" + "".join(f"{ran[e]:>13}" for e in engines))
    out.append("citation share = share of prompts where the domain appeared in the")
    out.append("engine's own cited sources. Not a ranking, and not a traffic estimate.")

    aud = r.get("audience") or {}
    if aud.get("engines"):
        out.append("")
        out.append(f"engine order and weights, {aud.get('slug')} "
                   f"(SparkToro report {aud.get('report_id')})")
        for name, w in sorted(aud["engines"].items(),
                              key=lambda kv: -kv[1]["audience_share"]):
            out.append(f"  {name:<12}{w['audience_share']:>6.1f}% of this buyer, "
                       f"{w['lift_pct']:+.0f}% vs the US average")

    weighted = audience_weighted(r)
    if weighted is not None:
        out.append("")
        out.append(f"reach-weighted citation share: {weighted:.0%}")
        out.append(f"  each engine weighted by the share of this buyer that uses it")
        out.append(f"  (SparkToro report {aud.get('report_id')}, audience {aud.get('slug')}).")
        out.append("  Covers the metered engines only. AI Overviews sits inside Google")
        out.append("  and has no citation API, so it is not in this number.")
    return "\n".join(out)


def pooled_share(r: dict, brand: str) -> float | None:
    """Share of (prompt x engine) pairs where this brand was in the sources.

    Pooled across engines rather than averaged, so an engine that answered more
    prompts carries more weight. Not a ranking and not a traffic estimate.
    """
    dom = r["domain"]
    hits = pairs = 0
    for e in r["engines"].values():
        n = e.get("prompts_run") or 0
        if not n:
            continue
        share = e["citation_share"] if brand == dom else e["competitor_share"].get(brand)
        if share is None:
            continue
        hits += share * n
        pairs += n
    return hits / pairs if pairs else None


def compare_arms(by_arm: dict) -> str:
    """The framing test: same buyer, same criteria, only the category word moves.

    Read it in this order, and no other:

    1. `soar` vs `automation` is the clean comparison. Neither arm mentions n8n,
       so if n8n.io shows up in one and not the other, the category label alone
       moved the competitive set.
    2. In the `n8n` arm every prompt names n8n, so n8n's own number there is
       primed and means nothing. The number that matters in that arm is
       tines.com: when buyers ask in the language the conquest book is buying,
       does Tines get cited at all?
    """
    arms = list(by_arm)
    if not arms:
        return ""
    first = by_arm[arms[0]]
    brands = [first["domain"]] + first["competitors"]
    bw = max(len(b) for b in brands) + 1

    out = ["\n\nframing comparison", "=" * (bw + 14 * len(arms)),
           f"{'brand':<{bw}}" + "".join(f"{a:>14}" for a in arms)]
    for b in brands:
        cells = []
        for a in arms:
            v = pooled_share(by_arm[a], b)
            cells.append("n/a" if v is None else f"{v:.0%}")
        mark = "*" if b == first["domain"] else " "
        out.append(f"{mark}{b:<{bw - 1}}" + "".join(f"{c:>14}" for c in cells))
    out.append("=" * (bw + 14 * len(arms)))
    out.append("pooled share of (prompt x engine) pairs citing that domain.")
    out.append("")
    out.append("prompt sets, one per arm, so each is its own series:")
    for a in arms:
        r = by_arm[a]
        out.append(f"  {a:<12} {r['prompt_count']:>2} prompts   set {r['prompt_set']}")
    out.append("")
    out.append("HOW TO READ THIS")
    out.append("  Arms that name no competitor are the clean comparison: a brand that")
    out.append("  moves between them moved because of the framing, nothing else.")
    primed = {a: sorted(b for b in brands
                        if b != first["domain"]
                        and all(mentions(p, b) for p in by_arm[a].get("_prompts", [])))
              for a in arms}
    for a, names in primed.items():
        if names:
            out.append(f"  Arm '{a}' names {', '.join(names)} in every prompt, so those")
            out.append(f"  numbers are primed and uninformative. Read {first['domain']} there.")
    return "\n".join(out)


def mentions(prompt: str, domain: str) -> bool:
    """Does this prompt name the brand behind a domain? Rough, and only used to
    warn that an arm has primed a competitor, never to score anything."""
    stem = domain.split(".")[0].lower()
    return len(stem) > 2 and stem in prompt.lower()


def audience_weighted(r: dict) -> float | None:
    """One number: citation share across engines, weighted by who uses them.

    Four separate engine percentages do not tell you how much of the buyer's AI
    surface you hold. This does, over the engines that actually returned data.
    """
    engines = (r.get("audience") or {}).get("engines") or {}
    if not engines:
        return None
    num = den = 0.0
    for name, e in r["engines"].items():
        share, w = e.get("citation_share"), engines.get(name, {}).get("audience_share")
        if share is None or not w:
            continue
        num += w * share
        den += w
    return num / den if den else None


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", required=True, help="JSON config: domain, competitors, prompts")
    ap.add_argument("--dry-run", action="store_true", help="print the shape without calling any engine")
    ap.add_argument("--trend", action="store_true", help="read the stored history and exit, no engine calls")
    ap.add_argument("--history", default="runs", help="directory of past runs (default: runs/)")
    ap.add_argument("--arm", help="run one named arm from a config that defines arms")
    ap.add_argument("--all-arms", action="store_true", help="run every arm and compare them")
    a = ap.parse_args()

    cfg = json.loads(open(a.config).read())
    if "domain" not in cfg:
        sys.exit("config is missing 'domain'")
    domain = cfg["domain"]
    runs_dir = Path(a.history)
    arms = cfg.get("arms") or {}

    if a.arm and a.arm not in arms:
        sys.exit(f"no arm {a.arm!r} in this config. available: {', '.join(arms) or 'none'}")
    if arms and not (a.arm or a.all_arms or a.trend):
        sys.exit(f"this config defines arms ({', '.join(arms)}). Pass --arm NAME for one, "
                 "or --all-arms to run them all and compare.")
    if not arms and "prompts" not in cfg:
        sys.exit("config is missing 'prompts'")

    # Which arms this invocation covers. None means the flat single-arm config.
    todo = list(arms) if a.all_arms else [a.arm] if a.arm else [None]

    # --trend never spends a token: it just reads what previous runs left behind.
    if a.trend:
        for arm in todo:
            label = f"{domain} [{arm}]" if arm else domain
            print(trend_table(load_history(runs_dir, domain, arm), label))
        return 0

    by_arm = {}
    for arm in todo:
        prompts = arms[arm]["prompts"] if arm else None
        if arm:
            print(f"\n=== arm: {arm} - {arms[arm].get('label', '')}", file=sys.stderr)
        res = run(cfg, a.dry_run, prompts=prompts, arm=arm)
        print(summarise(res))
        by_arm[arm or "all"] = res

        if a.dry_run:
            continue
        runs_dir.mkdir(parents=True, exist_ok=True)
        stamp = res["run_at"].replace(":", "").replace("-", "")
        out = run_path(runs_dir, domain, arm, stamp)
        out.write_text(json.dumps(res, indent=2))
        print(f"\nraw responses and per-prompt URLs written to {out}")

        # History includes the run just written, so this is the full series.
        history = load_history(runs_dir, domain, arm)
        if len(history) > 1:
            print(trend_table(history, f"{domain} [{arm}]" if arm else domain))
        else:
            print("\nfirst run stored. run it again next week and this becomes a trend.")

    if len(by_arm) > 1:
        print(compare_arms(by_arm))

    if a.dry_run:
        print("\ndry run: nothing written, nothing added to the series.")
    return 0


if __name__ == "__main__":
    sys.exit(main())
