"""adamo-bridge.py — local FastAPI bridge from a browser to the Adamo network.

Run this on the same laptop as your robot:

    pip install adamo fastapi uvicorn
    export ADAMO_API_KEY=ak_...
    python adamo-bridge.py           # binds http://localhost:7878

Then any Lovable app in the Physical AI creative-robotics repo can call:

    POST /put         { key, value }        -> session.put(key, value)
    GET  /subscribe?key=<expr>              -> SSE stream of samples
    GET  /alive                             -> array of live robot names
    POST /record/start { keys: [...] }      -> begin capturing samples
    POST /record/stop                       -> stop and return the take JSON
    POST /replay      { take, speed? }      -> re-publish the take at speed

Only the bridge holds ADAMO_API_KEY. The browser never sees it.

Built during the Physical AI & Creative Robotics Hackathon organised by
StreetKode Fam during Indian Krump Festival 14.
"""
from __future__ import annotations

import asyncio
import json
import os
import time
from typing import Any

import adamo  # pip install adamo
from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

API_KEY = os.environ.get("ADAMO_API_KEY")
if not API_KEY:
    raise SystemExit("Set ADAMO_API_KEY before running adamo-bridge.py")

session = adamo.connect(api_key=API_KEY)

app = FastAPI(title="adamo-bridge", version="1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # localhost bridge — a Lovable preview on any origin can call it
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------------------------------------------------------------------------
# /put

class PutBody(BaseModel):
    key: str
    value: Any


@app.post("/put")
def put(body: PutBody) -> dict[str, str]:
    payload = json.dumps(body.value).encode("utf-8")
    session.put(body.key, payload)
    return {"ok": "true"}


# ---------------------------------------------------------------------------
# /alive — declare a discovery watcher and return the current live names

_watched: dict[str, float] = {}


def _on_liveliness(sample: Any) -> None:
    key = getattr(sample, "key", None) or getattr(sample, "key_expr", None) or str(sample)
    _watched[str(key)] = time.time()


session.declare_subscriber("**/alive", _on_liveliness)


@app.get("/alive")
def alive() -> list[str]:
    cutoff = time.time() - 5
    return sorted({k.split("/")[0] for k, ts in _watched.items() if ts >= cutoff})


# ---------------------------------------------------------------------------
# /subscribe — SSE stream of samples for any key expression

@app.get("/subscribe")
async def subscribe(key: str = Query(...)) -> StreamingResponse:
    loop = asyncio.get_event_loop()
    q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    def _on_sample(sample: Any) -> None:
        try:
            raw = bytes(getattr(sample, "payload", b"") or b"")
            try:
                value = json.loads(raw.decode("utf-8")) if raw else None
            except Exception:
                value = raw.decode("utf-8", errors="replace")
            k = getattr(sample, "key", None) or getattr(sample, "key_expr", None) or key
            loop.call_soon_threadsafe(q.put_nowait, {"key": str(k), "value": value, "at": time.time()})
        except Exception:
            pass

    sub = session.declare_subscriber(key, _on_sample)

    async def event_stream():
        try:
            while True:
                sample = await q.get()
                yield f"data: {json.dumps(sample)}\n\n"
        finally:
            try:
                sub.undeclare()
            except Exception:
                pass

    return StreamingResponse(event_stream(), media_type="text/event-stream")


# ---------------------------------------------------------------------------
# /record + /stop + /replay

_recording: dict[str, Any] = {"active": False, "samples": [], "subs": [], "started_at": 0.0}


class RecordBody(BaseModel):
    keys: list[str]


@app.post("/record/start")
def record_start(body: RecordBody) -> dict[str, str]:
    if _recording["active"]:
        return {"ok": "already"}
    _recording["active"] = True
    _recording["samples"] = []
    _recording["started_at"] = time.time()

    def _cap(sample: Any) -> None:
        if not _recording["active"]:
            return
        try:
            raw = bytes(getattr(sample, "payload", b"") or b"")
            value = json.loads(raw.decode("utf-8")) if raw else None
        except Exception:
            value = None
        k = getattr(sample, "key", None) or getattr(sample, "key_expr", None) or ""
        _recording["samples"].append({
            "key": str(k),
            "value": value,
            "dt": time.time() - _recording["started_at"],
        })

    _recording["subs"] = [session.declare_subscriber(k, _cap) for k in body.keys]
    return {"ok": "true"}


@app.post("/record/stop")
def record_stop() -> dict[str, Any]:
    if not _recording["active"]:
        return {"samples": []}
    _recording["active"] = False
    for s in _recording["subs"]:
        try:
            s.undeclare()
        except Exception:
            pass
    _recording["subs"] = []
    return {"samples": list(_recording["samples"]), "duration": time.time() - _recording["started_at"]}


class ReplayBody(BaseModel):
    take: dict[str, Any]
    speed: float = 1.0


@app.post("/replay")
async def replay(body: ReplayBody) -> dict[str, str]:
    samples = body.take.get("samples", [])
    if not samples:
        return {"ok": "empty"}
    start = time.time()
    speed = max(0.1, body.speed)
    for s in samples:
        target = start + (s["dt"] / speed)
        delay = target - time.time()
        if delay > 0:
            await asyncio.sleep(delay)
        session.put(s["key"], json.dumps(s.get("value")).encode("utf-8"))
    return {"ok": "true"}


# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=7878, log_level="info")