#!/usr/bin/env python3
"""Lab 2: the same queries through Exa, Firecrawl and Tavily search APIs.

Records per call: latency, result count, result URLs, provider-reported cost
or credits, and errors. Output: results/<run-id>/<provider>.jsonl.

Keys come from environment variables EXA_API_KEY, FIRECRAWL_API_KEY and
TAVILY_API_KEY; a provider without a key is skipped.

Request shapes:
- Exa: the build-with-exa recommended /search request (query, type auto,
  contents.highlights) and nothing else. Exa returns 10 results by default.
- Firecrawl /v2/search and Tavily /search default to 5 results, so both are
  set to 10 to match Exa's default (an intentional comparison decision).

usage: run.py queries.txt [--providers exa,firecrawl,tavily] [--limit N]
"""
import json, os, sys, time, urllib.request, urllib.error, datetime, argparse

CRED = "/opt/pondero.ai/ops/credentials/test-accounts.env"

def load_keys():
    """API keys from the environment (EXA_API_KEY, TAVILY_API_KEY,
    FIRECRAWL_API_KEY), plus PONDERO_CRED_FILE / the ops-box file if present."""
    keys = {k: v for k, v in os.environ.items() if k.endswith("_API_KEY")}
    path = os.environ.get("PONDERO_CRED_FILE", CRED)
    if not os.path.exists(path):
        return keys
    for line in open(path):
        line = line.strip()
        if "=" in line and not line.startswith("#"):
            k, v = line.split("=", 1)
            keys[k] = v.strip().strip('"').strip("'")
    return keys

def post(url, headers, body, timeout=60):
    req = urllib.request.Request(url, json.dumps(body).encode(), {"Content-Type": "application/json", **headers})
    t0 = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            data = json.load(r); status = r.status
    except urllib.error.HTTPError as e:
        data = {"error": e.read().decode(errors="replace")[:500]}; status = e.code
    except Exception as e:
        data = {"error": str(e)}; status = 0
    return status, data, round((time.perf_counter() - t0) * 1000)

def exa(key, q):
    s, d, ms = post("https://api.exa.ai/search", {"x-api-key": key},
                    {"query": q, "type": "auto", "contents": {"highlights": True}})
    res = d.get("results", []) if s == 200 else []
    items = [{"url": r.get("url"), "title": r.get("title"), "snippet": " ".join(r.get("highlights") or [])[:400]} for r in res]
    return s, ms, items, {"costDollars": d.get("costDollars"), "searchTime": d.get("searchTime")}, d.get("error")

def firecrawl(key, q):
    s, d, ms = post("https://api.firecrawl.dev/v2/search", {"Authorization": "Bearer " + key},
                    {"query": q, "limit": 10})
    web = ((d.get("data") or {}).get("web") or []) if s == 200 else []
    items = [{"url": r.get("url"), "title": r.get("title"), "snippet": (r.get("description") or "")[:400]} for r in web]
    return s, ms, items, {"creditsUsed": d.get("creditsUsed")}, d.get("error")

def tavily(key, q):
    s, d, ms = post("https://api.tavily.com/search", {"Authorization": "Bearer " + key},
                    {"query": q, "max_results": 10})
    res = d.get("results", []) if s == 200 else []
    items = [{"url": r.get("url"), "title": r.get("title"), "snippet": (r.get("content") or "")[:400]} for r in res]
    return s, ms, items, {"responseTime": d.get("response_time"), "usage": d.get("usage")}, d.get("error") or d.get("detail")

PROVIDERS = {"exa": ("EXA_API_KEY", exa), "firecrawl": ("FIRECRAWL_API_KEY", firecrawl), "tavily": ("TAVILY_API_KEY", tavily)}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("queries"); ap.add_argument("--providers", default="exa,firecrawl,tavily")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--sleep", type=float, default=0, help="seconds between calls (Firecrawl free plan: 10 searches/min)")
    ap.add_argument("--out", default="", help="write into an existing results dir")
    a = ap.parse_args()
    qs = []
    for l in open(a.queries):
        l = l.strip()
        if not l or l.startswith("#"): continue
        cat, _, q = l.partition("|") if "|" in l else ("", "", l)
        qs.append((cat.strip(), q.strip()))
    if a.limit: qs = qs[:a.limit]
    keys = load_keys()
    run = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
    out = a.out or os.path.join(os.path.dirname(os.path.abspath(__file__)), "results", run); os.makedirs(out, exist_ok=True)
    for name in a.providers.split(","):
        env, fn = PROVIDERS[name]
        if not keys.get(env):
            print(f"skip {name}: no {env}"); continue
        with open(os.path.join(out, name + ".jsonl"), "w") as f:
            for cat, q in qs:
                s, ms, items, meta, err = fn(keys[env], q)
                f.write(json.dumps({"provider": name, "category": cat, "query": q, "status": s, "ms": ms, "n": len(items), "results": items, "meta": meta, "error": err}) + "\n")
                print(f"{name:9} {s} {ms:6}ms n={len(items):2} {q[:60]}")
                if a.sleep: time.sleep(a.sleep)
    print("results:", out)

if __name__ == "__main__":
    main()