#!/usr/bin/env python3
"""Cheap connectivity probe for providers.json entries.
Never prints the API key. Prefer --once / --max-probes / scheduled --minutes/--hours.

  python probe.py --mock
  python probe.py --provider flint --once
  python probe.py --provider flint --minutes 60 --hours 8
"""
from __future__ import annotations
import argparse, csv, json, os, sys, time

CSV_DEFAULT = "results/llm_probe.csv"


def _append_csv(path, row, fieldnames):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    new = not os.path.exists(path) or os.path.getsize(path) == 0
    with open(path, "a", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        if new:
            w.writeheader()
        w.writerow(row)


def _scrub(err, key):
    s = str(err)
    if key:
        s = s.replace(key, "[redacted]")
        if len(key) > 8:
            s = s.replace(key[:8], "[redacted]")
    return s


def main():
    ap = argparse.ArgumentParser(description="Probe an OpenAI-compatible provider without leaking secrets.")
    ap.add_argument("--config", default="providers.json")
    ap.add_argument("--provider", default="flint")
    ap.add_argument("--mock", action="store_true", help="offline plumbing check only")
    ap.add_argument("--once", action="store_true", help="single live call then exit")
    ap.add_argument("--max-probes", type=int, default=1, help="max live calls (default 1)")
    ap.add_argument("--minutes", type=float, default=0, help="spacing between probes in a long run")
    ap.add_argument("--hours", type=float, default=0, help="total wall time for a long run")
    ap.add_argument("--csv", default=CSV_DEFAULT, help="append probe outcomes here (never writes the key)")
    ap.add_argument("--prompt", default='Return JSON {"ok":true} and nothing else.')
    a = ap.parse_args()
    fields = ["ts_rel", "provider", "model", "status", "seconds", "error", "usage_calls"]

    if a.mock:
        from interfaces import MockLLMClient
        c = MockLLMClient([{"ok": True}])
        out = c.complete_json(a.prompt)
        print("mock_ok", bool(out.get("ok")), "calls", c.usage.get("calls"))
        return 0

    cfg = json.load(open(a.config))
    local = "fixtures/private_unredacted/providers.local.json"
    if os.path.exists(local):
        loc = json.load(open(local))
        if a.provider in loc:
            cfg.setdefault(a.provider, {}).update({k: v for k, v in loc[a.provider].items() if not str(k).startswith("_")})

    p = cfg.get(a.provider)
    if not p:
        print(f"provider_missing {a.provider}", file=sys.stderr); return 2
    env = p.get("env", "SQ_LLM_API_KEY")
    key = os.environ.get(env, "")
    if not key:
        print(f"key_missing env={env} provider={a.provider}"); return 3
    print(f"key_present env={env} len={len(key)} provider={a.provider}")

    from interfaces import OpenAICompatClient
    model = p.get("model") or os.environ.get(p.get("model_env", "SQ_LLM_MODEL"), "")
    base = p.get("base_url") or os.environ.get(p.get("base_url_env", "SQ_LLM_BASE_URL"), "")
    client = OpenAICompatClient(model=model, base_url=base, api_key=key,
                                timeout=p.get("timeout", 120), sleep_s=p.get("sleep_s", 0),
                                json_mode=p.get("json_mode", True))

    n = 1 if a.once else max(1, a.max_probes)
    end = time.time() + (a.hours * 3600 if a.hours else 0)
    i = 0
    status = "fail"
    while i < n or (a.hours and time.time() < end):
        i += 1
        t0 = time.time()
        try:
            out = client.complete_json(a.prompt)
            status = "ok" if out is not None else "empty"
            err = ""
        except Exception as e:
            status = "fail"; err = _scrub(e, key); out = None
        sec = time.time() - t0
        _append_csv(a.csv, {"ts_rel": f"probe-{i}", "provider": a.provider, "model": model,
                            "status": status, "seconds": f"{sec:.3f}", "error": err,
                            "usage_calls": client.usage.get("calls", 0)}, fields)
        print(f"probe {i} status={status} seconds={sec:.2f} error={err}")
        if a.once or (not a.hours and i >= n):
            break
        if a.minutes:
            time.sleep(a.minutes * 60)
    return 0 if status == "ok" else 1


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