#!/usr/bin/env python3
"""
Agent Hub — Local Bridge.

Run this on WHATEVER machine your agent lives on (laptop, VM, server behind a
firewall). It makes only OUTBOUND requests to the hub — no ports to open, no
tunnels — long-polling for jobs and streaming results back.

Two ways to wire your agent in:

  1. Run a command per job (job input arrives on stdin, stdout streams back):
       python3 bridge.py --url https://hub.example.com --token brg_xxx \
           --exec "python3 my_agent.py"

  2. Forward jobs to an agent already serving HTTP locally:
       python3 bridge.py --url https://hub.example.com --token brg_xxx \
           --forward http://localhost:8000/run
     The bridge POSTs {"run_id","input","user"} and accepts JSON
     {"output": "..."}, plain text, or a text/event-stream of data: chunks.

Only the Python 3 standard library is required.
"""
import argparse
import codecs
import json
import os
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.request
import urllib.error


def api(base, token, path, method="GET", body=None, timeout=40):
    req = urllib.request.Request(
        base.rstrip("/") + path,
        data=json.dumps(body).encode() if body is not None else None,
        method=method,
        headers={"Authorization": "Bearer " + token,
                 "Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def send_event(base, token, run_id, etype, data="", seq=None):
    """Deliver one event to the hub, retrying over transient network blips.

    Terminal events (done/error) get more attempts — losing one turns a
    finished answer into a bogus five-minute timeout on the hub. Chunk
    events carry a sequence number so a retry that follows an
    already-processed request is deduplicated hub-side.
    """
    body = {"type": etype, "data": data}
    if seq is not None:
        body["seq"] = seq
    attempts = 5 if etype in ("done", "error") else 3
    delay = 1
    for i in range(attempts):
        try:
            api(base, token, f"/api/bridge/runs/{run_id}/events",
                method="POST", body=body, timeout=30)
            return True
        except Exception as e:
            if i == attempts - 1:
                print(f"  ! gave up sending {etype} for {run_id}: {e}",
                      file=sys.stderr)
                return False
            time.sleep(delay)
            delay = min(delay * 2, 8)


def download_files(job):
    """Fetch the job's attachments into a temp dir; returns (dir, paths).

    Attachments arrive as capability URLs on the hub; the agent gets plain
    local file paths and never needs network access itself.
    """
    files = job.get("files") or []
    if not files:
        return None, [], []
    tmp = tempfile.mkdtemp(prefix="hubjob_")
    paths, failed = [], []
    for i, f in enumerate(files):
        base = os.path.basename(f.get("name") or "") or f"file_{i}"
        path = os.path.join(tmp, base)
        if os.path.exists(path):
            path = os.path.join(tmp, f"{i}_{base}")
        try:
            with urllib.request.urlopen(f["url"], timeout=120) as resp, \
                    open(path, "wb") as out:
                shutil.copyfileobj(resp, out)
            paths.append(path)
        except Exception as e:
            print(f"  ! could not download attachment "
                  f"{f.get('name')}: {e}", file=sys.stderr)
            failed.append(f.get("name") or f"file_{i}")
    return tmp, paths, failed


def job_stdin(job, file_paths=(), failed_files=()):
    """What an --exec agent reads on stdin.

    Single-turn jobs get the raw input. Replies in a conversation get the
    whole thread as a plain-text transcript, newest message last. Attached
    files are downloaded locally and listed by path at the end; downloads
    that failed are named explicitly so the agent (and its answer) can say
    so instead of silently working without them.
    """
    msgs = job.get("messages") or []
    if len(msgs) <= 1:
        text = job["input"]
    else:
        lines = ["[conversation so far]"]
        for m in msgs[:-1]:
            who = "User" if m.get("role") == "user" else "Agent"
            lines.append(f"{who}: {m.get('content', '')}")
        lines += ["", "[new message]", job["input"]]
        text = "\n".join(lines)
    if file_paths or failed_files:
        text += "\n\n[attached files]\n" + "\n".join(file_paths)
        for name in failed_files:
            text += f"\nFAILED: {name} (the attachment could not be downloaded)"
    return text


def handle_exec(args, job):
    run_id = job["run_id"]
    print(f"→ job {run_id}: running `{args.exec}`")
    tmp_dir, file_paths, failed_files = download_files(job)
    try:
        proc = subprocess.Popen(
            args.exec, shell=True,
            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    except Exception as e:
        send_event(args.url, args.token, run_id, "error", f"Could not start agent: {e}")
        if tmp_dir:
            shutil.rmtree(tmp_dir, ignore_errors=True)
        return

    def feed():
        try:
            proc.stdin.write(
                job_stdin(job, file_paths, failed_files).encode("utf-8"))
            proc.stdin.close()
        except Exception:
            pass

    threading.Thread(target=feed, daemon=True).start()

    # drain stderr concurrently — an agent that logs more than the ~64KB
    # pipe buffer would otherwise deadlock against our stdout read
    err_tail = []

    def drain_stderr():
        try:
            while True:
                b = os.read(proc.stderr.fileno(), 65536)
                if not b:
                    return
                err_tail.append(b)
                del err_tail[:-8]           # keep a bounded tail
        except OSError:
            pass

    err_thread = threading.Thread(target=drain_stderr, daemon=True)
    err_thread.start()

    # os.read returns whatever is available (true live streaming), and the
    # incremental decoder keeps multibyte characters split across chunk
    # boundaries intact
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    sent_any, seq, lost = False, 0, 0

    def send_chunk(text):
        nonlocal sent_any, seq, lost
        sent_any = True
        seq += 1
        if not send_event(args.url, args.token, run_id, "chunk", text, seq=seq):
            lost += 1

    while True:
        chunk = os.read(proc.stdout.fileno(), 65536)
        if not chunk:
            break
        text = decoder.decode(chunk)
        if text:
            send_chunk(text)
    tail = decoder.decode(b"", final=True)
    if tail:
        send_chunk(tail)
    proc.wait()
    err_thread.join(timeout=3)  # let the drain catch the final stderr writes
    if lost:
        seq += 1
        send_event(args.url, args.token, run_id, "chunk",
                   f"\n[{lost} piece(s) of output were lost in transmission]\n",
                   seq=seq)
    if proc.returncode != 0:
        err = b"".join(err_tail).decode("utf-8", "replace")[-2000:]
        send_event(args.url, args.token, run_id, "error",
                   err or f"Agent exited with code {proc.returncode}.")
    else:
        send_event(args.url, args.token, run_id, "done",
                   "" if sent_any else "(the agent produced no output)")
    if tmp_dir:
        shutil.rmtree(tmp_dir, ignore_errors=True)
    print(f"✓ job {run_id} finished (exit {proc.returncode})")


def handle_forward(args, job):
    run_id = job["run_id"]
    print(f"→ job {run_id}: forwarding to {args.forward}")
    req = urllib.request.Request(
        args.forward, method="POST",
        data=json.dumps(job).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            ctype = resp.headers.get("Content-Type", "")
            if "text/event-stream" in ctype:
                for raw in resp:
                    line = raw.decode("utf-8", "replace").rstrip("\r\n")
                    if line.startswith("data:"):
                        send_event(args.url, args.token, run_id, "chunk",
                                   line[5:].lstrip() + "\n")
                send_event(args.url, args.token, run_id, "done")
            else:
                body = resp.read().decode("utf-8", "replace")
                try:
                    parsed = json.loads(body)
                    if isinstance(parsed, dict) and "output" in parsed:
                        body = str(parsed["output"])
                except ValueError:
                    pass
                send_event(args.url, args.token, run_id, "done", body)
        print(f"✓ job {run_id} finished")
    except Exception as e:
        send_event(args.url, args.token, run_id, "error",
                   f"Local agent unreachable: {e}")


def main():
    ap = argparse.ArgumentParser(description="Agent Hub local bridge")
    ap.add_argument("--url", required=True, help="Hub base URL, e.g. https://hub.example.com")
    ap.add_argument("--token", required=True, help="Bridge token from the Connect agent flow")
    mode = ap.add_mutually_exclusive_group(required=True)
    mode.add_argument("--exec", help="Shell command run per job (input on stdin, stdout streamed back)")
    mode.add_argument("--forward", help="Local HTTP endpoint to POST jobs to, e.g. http://localhost:8000/run")
    args = ap.parse_args()

    print(f"Bridge connected to {args.url} — waiting for jobs (Ctrl-C to stop).")
    backoff = 2
    while True:
        try:
            data = api(args.url, args.token, "/api/bridge/jobs", timeout=40)
            backoff = 2
            job = data.get("job")
            if job:
                handler = handle_exec if args.exec else handle_forward
                threading.Thread(target=handler, args=(args, job), daemon=True).start()
        except KeyboardInterrupt:
            print("\nBridge stopped.")
            return
        except urllib.error.HTTPError as e:
            if e.code == 401:
                print("Hub rejected the token (was it rotated?). Exiting.", file=sys.stderr)
                return
            print(f"Hub error {e.code}; retrying in {backoff}s", file=sys.stderr)
            time.sleep(backoff)
            backoff = min(backoff * 2, 60)
        except Exception as e:
            print(f"Connection problem ({e}); retrying in {backoff}s", file=sys.stderr)
            time.sleep(backoff)
            backoff = min(backoff * 2, 60)


if __name__ == "__main__":
    main()