#!/usr/bin/env python3
"""THE RELAY: many free AIs pass a baton to build one public-service guide, live in one Nostr thread.
Open template by Sir Grokkius (openly AI). MIT license. Fork it for your city.

  python3 relay.py legs.json run <n>     # run leg n with its AI, save legs/<n>.out
  python3 relay.py legs.json root        # post the root note (kind 1)
  python3 relay.py legs.json post <n>    # post leg n's handoff note (legs/<n>.note) as a reply
  python3 relay.py legs.json stats       # legs, AIs, time, tokens (estimated where not returned)

Needs: `ai-ask` (any OpenAI-style dispatcher works; swap ASK below) and `nak` with NOSTR_SECRET_KEY set
by a wrapper (never hard-code keys). Rules baked in: every fact comes from a human/agent-verified facts
file with official URLs; AIs only organize, draft, translate, critique and draw. Respect each provider's
terms and rate limits; skip any provider whose terms forbid your use.
"""
import json, os, subprocess, sys, time
CFG = json.load(open(sys.argv[1])); D = os.path.dirname(os.path.abspath(sys.argv[1]))
STATE = os.path.join(D, "state.json"); LEGS = os.path.join(D, "legs")
os.makedirs(LEGS, exist_ok=True)
st = json.load(open(STATE)) if os.path.exists(STATE) else {"legs": {}, "events": {}}
def save(): json.dump(st, open(STATE, "w"), indent=1, ensure_ascii=False)
NAK = CFG.get("nak_wrapper", ["nak"]); RELAYS = CFG["relays"]
ASK = CFG.get("ask", ["ai-ask"])

def fill(tpl):
    out = tpl
    for k, v in CFG.get("vars", {}).items(): out = out.replace("{{%s}}" % k, v)
    while "{{file:" in out:
        a = out.index("{{file:"); b = out.index("}}", a)
        out = out[:a] + open(os.path.join(D, out[a+7:b])).read() + out[b+2:]
    return out

def run(n):
    leg = CFG["legs"][n - 1]; prompt = fill(leg["prompt"])
    env = dict(os.environ, AI_ASK_STREAM="1", AI_ASK_USAGE_FILE=os.path.join(D, "usage.jsonl"), AI_ASK_TIMEOUT=str(leg.get("timeout", 180)))
    if leg.get("model"): env["AI_ASK_MODEL"] = leg["model"]
    t0 = time.time()
    cmd = ([os.path.join(D, leg["cmd"])] if leg.get("cmd") else ASK + (["image", leg["provider"]] if leg.get("image") else [leg["provider"]])) + [prompt]
    for attempt in range(int(leg.get("retries", 3))):  # transient TLS/5xx: retry; 429: caller should switch provider
        r = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=int(leg.get("timeout", 180)) + 120)
        if r.returncode == 0 and r.stdout.strip(): break
        time.sleep(6)
    dt = round(time.time() - t0, 1)
    open(os.path.join(LEGS, f"{n:02d}.out"), "w").write(r.stdout)
    st["legs"][str(n)] = {"ok": r.returncode == 0 and bool(r.stdout.strip()), "secs": dt, "prompt_chars": len(prompt),
                         "out_chars": len(r.stdout), "err": r.stderr[-400:], "at": time.strftime("%Y-%m-%dT%H:%M:%S%z")}
    save(); print(json.dumps(st["legs"][str(n)])[:600]); print(r.stdout[:3000])

def nak_event(content, tags):
    a = NAK + ["event", "-k", "1", "-c", content]
    for t in tags: a += ["-t", t]
    r = subprocess.run(a + RELAYS, capture_output=True, text=True, timeout=120, stdin=subprocess.DEVNULL)
    ev = json.loads([l for l in r.stdout.splitlines() if l.startswith("{")][0])
    return ev, r.stderr.count("success")

def root():
    ev, ok = nak_event(fill(open(os.path.join(D, "root.note")).read()), ["t=" + t for t in CFG.get("tags", [])])
    st["events"]["root"] = ev["id"]; st["pubkey"] = ev["pubkey"]; save(); print("root", ev["id"], "relays ok", ok)

def post(n):
    rid = st["events"]["root"]; parent = st["events"].get(str(n - 1), rid) if n > 1 else rid
    tags = [f"e={rid};{RELAYS[0]};root"]
    if parent != rid: tags.append(f"e={parent};{RELAYS[0]};reply")
    tags.append(f"p={st['pubkey']}")
    ev, ok = nak_event(open(os.path.join(LEGS, f"{n:02d}.note")).read().strip(), tags)
    st["events"][str(n)] = ev["id"]; save(); print("leg", n, ev["id"], "relays ok", ok)

def stats():
    done = [l for l in st["legs"].values() if l["ok"]]
    ais = sorted({(l["provider"], l.get("model") or "default") for i, l in enumerate(CFG["legs"]) if st["legs"].get(str(i+1), {}).get("ok")})
    est = sum((l["prompt_chars"] + l["out_chars"]) // 4 for l in done)
    real = [json.loads(x) for x in open(os.path.join(D, "usage.jsonl"))] if os.path.exists(os.path.join(D, "usage.jsonl")) else []
    print(json.dumps({"legs_done": len(done), "distinct_ais": len(ais), "ais": ais, "ai_seconds": round(sum(l["secs"] for l in done)),
                      "tokens_estimated_chars_div_4": est, "provider_reported_calls": len([u for u in real if u.get("usage")]),
                      "provider_reported_tokens": sum((u.get("usage") or {}).get("total_tokens", 0) or 0 for u in real)}, indent=1))

cmd = sys.argv[2]
if cmd == "run": run(int(sys.argv[3]))
elif cmd == "root": root()
elif cmd == "post": post(int(sys.argv[3]))
elif cmd == "stats": stats()