#!/usr/bin/env python3
"""
sparktoro.py, audience-side input for the citation meter.

citation_meter.py answers "who gets cited". It cannot answer the two questions
that come before that, and for a while I was answering both by guessing:

  1. Which engines should we meter for THIS buyer?
  2. Which questions should we meter them on?

SparkToro answers both from clickstream panel data instead of intuition. This
script pulls two endpoints per audience and writes them next to the config so
every number on a microsite traces back to a file on disk.

  /v3/apps/ai   which AI assistants the audience actually uses, as a share of
                the audience, with the national baseline alongside it. That is
                the engine list and its weighting.
  /v3/prompts   the AI prompt topics the audience carries affinity for. That is
                the topic set the prompts get written against.

One thing this does NOT do: SparkToro reports topics, not verbatim prompts, and
it reports assistant USAGE, not citations. It tells us where the buyer is and
what they ask about. Only the meter can tell us whether we are named there. The
suggested prompts below are templates over the topics, marked as derived, and a
human rewrites them before they count as a measurement.

Usage
  export SPARKTORO_API_KEY=...
  ./sparktoro.py --describe "security operations engineers who evaluate SOAR"
                 --slug tines-com
  ./sparktoro.py --slug tines-com --report                 # read the cache
  ./sparktoro.py --slug tines-com --sync-config tines.json # merge engine weights

Credits (a report is 10, prompts 2, ai apps 2, so a fresh audience is 14):
  ./sparktoro.py --credits
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

API = "https://api.sparktoro.com"

# Cache lives next to the CONFIGS, not next to this script. That distinction
# matters: this file is tracked and published, and an audience pull names the
# company it is about. Resolving the cache script-relative would drop private
# reads into a public directory. Default to the working directory instead, and
# let --cache move it.
CACHE = Path("audience")

# SparkToro reports assistants by domain. The meter knows them by engine name.
# Anything not in this map is still reported, it just is not directly meterable
# (google.com is AI Overviews, which the meter reaches through the SERP, not an
# API of its own).
ENGINE_BY_DOMAIN = {
    "chatgpt.com": "openai",
    "claude.ai": "anthropic",
    "gemini.google.com": "gemini",
    "perplexity.ai": "perplexity",
}

# A country_average of 0 means SparkToro has no national baseline for that
# domain, so its "lift" comes back as 999999. Those rows are noise and are never
# quoted. This is the floor a baseline has to clear to be usable.
MIN_BASELINE = 0.5


def curl(path: str, key: str, post: dict | None = None) -> dict:
    """The sandbox proxy is configured for curl, not urllib. Shell out."""
    cmd = ["curl", "-sS", "-m", "90", f"{API}{path}",
           "-H", f"Authorization: Bearer {key}"]
    if post is not None:
        cmd += ["-X", "POST", "-H", "Content-Type: application/json",
                "-d", json.dumps(post)]
    out = subprocess.run(cmd, capture_output=True, text=True, check=True).stdout
    try:
        return json.loads(out)
    except json.JSONDecodeError:
        raise SystemExit(f"sparktoro: unreadable response from {path}: {out[:300]}")


def poll(path: str, key: str, tries: int = 12, wait: int = 10) -> dict:
    """Report sections return 202 while they build. Wait them out."""
    for attempt in range(tries):
        d = curl(path, key)
        if not d.get("error"):
            return d
        if d.get("status") != 202:
            raise SystemExit(f"sparktoro: {d.get('status')} {d.get('message')}")
        print(f"  preparing, retry {attempt + 1}/{tries}", file=sys.stderr)
        time.sleep(wait)
    raise SystemExit("sparktoro: report never finished building")


def credits(key: str) -> dict:
    return curl("/v3/account/credits", key)


def describe(prompt: str, key: str) -> str:
    d = curl("/v3/describe/create", key, post={"prompt": prompt})
    if d.get("error"):
        raise SystemExit(f"sparktoro: {d.get('message')}")
    return d["report_id"]


def pull(slug: str, report_id: str, key: str) -> tuple[dict, dict]:
    CACHE.mkdir(parents=True, exist_ok=True)
    prompts = poll(f"/v3/prompts?report_id={report_id}", key)
    apps = poll(f"/v3/apps/ai?report_id={report_id}", key)
    for name, blob in (("prompts", prompts), ("ai-apps", apps)):
        blob["_report_id"] = report_id
        (CACHE / f"{slug}-{name}.json").write_text(json.dumps(blob, indent=1))
    return prompts, apps


def load(slug: str) -> tuple[dict, dict]:
    p = CACHE / f"{slug}-prompts.json"
    a = CACHE / f"{slug}-ai-apps.json"
    if not p.exists() or not a.exists():
        raise SystemExit(f"sparktoro: no cache for {slug}. Run --describe first.")
    return json.loads(p.read_text()), json.loads(a.read_text())


def assistants(apps: dict) -> list[dict]:
    """Rows with a real national baseline, richest lift first."""
    rows = [r for r in apps["data"]
            if (r.get("country_average") or 0) >= MIN_BASELINE]
    return sorted(rows, key=lambda r: -r["affinity"])


def engine_weights(apps: dict) -> dict:
    """Meterable engines, ordered by how much this audience over-indexes."""
    out = {}
    for r in assistants(apps):
        engine = ENGINE_BY_DOMAIN.get(r["domain"])
        if engine:
            out[engine] = {
                "domain": r["domain"],
                "audience_share": round(r["affinity"], 2),
                "us_average": round(r["country_average"], 2),
                "lift_pct": round(r["affinity_lift"], 1),
            }
    return out


def suggested_prompts(prompts: dict, top: int = 6) -> list[str]:
    """Templates over the top topics. Drafts, not measurements. Rewrite them."""
    rows = sorted(prompts["data"], key=lambda r: -r["affinity"])[:top]
    return [f"What are the best tools for {r['topic'].lower()}?" for r in rows]


def report(slug: str) -> str:
    prompts, apps = load(slug)
    rows = assistants(apps)
    lines = [f"AUDIENCE: {slug}", ""]

    lines.append("Assistants this audience uses (SparkToro /v3/apps/ai)")
    lines.append(f"  {'assistant':24} {'audience':>9} {'US avg':>9} {'lift':>9}  meter")
    for r in rows[:12]:
        engine = ENGINE_BY_DOMAIN.get(r["domain"], "")
        lines.append(f"  {r['domain']:24} {r['affinity']:8.2f}% "
                     f"{r['country_average']:8.2f}% {r['affinity_lift']:8.1f}%  {engine}")

    lines += ["", "AI prompt topics by affinity (SparkToro /v3/prompts)"]
    for r in sorted(prompts["data"], key=lambda x: -x["affinity"])[:12]:
        lines.append(f"  {r['affinity']:5.0f}  {r['topic']}")

    lines += ["", "Meterable engines, by audience weight"]
    for engine, w in engine_weights(apps).items():
        lines.append(f"  {engine:12} {w['audience_share']:6.2f}% of audience, "
                     f"{w['lift_pct']:+.0f}% vs US")

    lines += ["", "Suggested prompts, DERIVED FROM TOPICS, rewrite before use"]
    lines += [f"  - {p}" for p in suggested_prompts(prompts)]
    return "\n".join(lines)


def sync_config(slug: str, config_path: Path) -> None:
    """Merge the audience read into a citation_meter config, non-destructively.

    Hand-written prompts are never overwritten. The audience block is added
    alongside them so the meter's engine order and the page's claims share one
    source, and so a stale audience read is visible rather than assumed.
    """
    prompts, apps = load(slug)
    cfg = json.loads(config_path.read_text()) if config_path.exists() else {}
    cfg["audience"] = {
        "source": "sparktoro",
        "slug": slug,
        "report_id": apps.get("_report_id"),
        "engines": engine_weights(apps),
        "topics": [r["topic"] for r in
                   sorted(prompts["data"], key=lambda x: -x["affinity"])[:10]],
    }
    cfg.setdefault("prompts", suggested_prompts(prompts))
    config_path.write_text(json.dumps(cfg, indent=2) + "\n")
    print(f"merged audience {slug} into {config_path.name}, "
          f"{len(cfg['audience']['engines'])} meterable engines")


def main() -> int:
    global CACHE
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--slug", help="cache name, e.g. tines-com")
    ap.add_argument("--describe", help="audience description, spends 14 credits")
    ap.add_argument("--report-id", help="reuse an existing SparkToro report")
    ap.add_argument("--report", action="store_true", help="print the cached read")
    ap.add_argument("--sync-config", type=Path, help="merge into a meter config")
    ap.add_argument("--credits", action="store_true", help="check the balance")
    ap.add_argument("--cache", type=Path, default=CACHE,
                    help="where audience pulls are stored (default: ./audience)")
    args = ap.parse_args()
    CACHE = args.cache

    key = os.environ.get("SPARKTORO_API_KEY", "")
    needs_key = args.describe or args.report_id or args.credits
    if needs_key and not key:
        raise SystemExit("sparktoro: set SPARKTORO_API_KEY")

    if args.credits:
        c = credits(key)
        print(f"{c['credits_remaining']} credits, expires {c['credits_expires_at']}"
              f"{', LOW' if c.get('low_balance') else ''}")
        return 0

    if args.describe or args.report_id:
        if not args.slug:
            raise SystemExit("sparktoro: --slug is required when pulling")
        rid = args.report_id or describe(args.describe, key)
        print(f"report {rid}", file=sys.stderr)
        pull(args.slug, rid, key)
        print(f"cached to {CACHE}/{args.slug}-*.json", file=sys.stderr)

    if args.report:
        print(report(args.slug))
    if args.sync_config:
        sync_config(args.slug, args.sync_config)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
