"""
SKYKOI koi-brain — the koi-harness multi-agent orchestrator as a device service.

WHY THIS EXISTS
---------------
The koi's slow lane is one model improvising with tools inside a vendored
coding-agent loop. koi-harness showed the same model doing better work when a
router gates the cognition depth, a planner fans independent steps out in
parallel, shared facts are derived once and broadcast, and the final answer
passes an invariant audit before anyone sees it. This service is that
orchestrator, vendored (./koibrain), reachable the same way the browser agent
is: loopback HTTP, NDJSON progress, one warm process.

Two deliberate seams:
- Browser steps are DELEGATED to apps/browser-agent (:18793) — one browser,
  one profile, one set of logins on the machine (koibrain/browser_worker.py).
- Sessions are isolated: each sessionId gets its own Orchestrator (its own
  conversation history), all sharing one SQLite memory under ~/.skykoi/koi-brain.

Run:
    python server.py            # binds 127.0.0.1:18794

POST /run {"task": "...", "sessionId": "voice"}  →  NDJSON: route / plan /
step / step_done / done{result}.
"""

from __future__ import annotations

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


# ---- self-bootstrap: install our python deps on first run -------------------
# A fresh machine has Python but not our packages (the koi installer provisions
# Node; python deps land here on demand). Missing imports: pip install from the
# requirements.txt beside this file, then re-exec once. The env flag guards
# against an install-fail loop; exit 78 tells the supervisor "configuration".
def _ensure_python_deps(probe_modules) -> None:
    import importlib.util
    import os as _os
    import subprocess as _sp
    import sys as _sys

    if all(importlib.util.find_spec(m) is not None for m in probe_modules):
        return
    if _os.environ.get("SKYKOI_PY_BOOTSTRAPPED") == "1":
        print("[bootstrap] python deps still missing after install; giving up", flush=True)
        _sys.exit(78)
    req = _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "requirements.txt")
    print(f"[bootstrap] installing python deps from {req} (first run)…", flush=True)
    for cmd in (
        [_sys.executable, "-m", "pip", "install", "-r", req],
        [_sys.executable, "-m", "pip", "install", "--user", "-r", req],
    ):
        try:
            if _sp.run(cmd, timeout=900).returncode == 0:
                break
        except Exception:
            continue
    env = {**_os.environ, "SKYKOI_PY_BOOTSTRAPPED": "1"}
    _os.execve(_sys.executable, [_sys.executable, _os.path.abspath(__file__)], env)

_ensure_python_deps(("aiohttp", "openai", "httpx", "dotenv", "rich"))

from aiohttp import web

HOST = os.environ.get("SKYKOI_KOI_BRAIN_HOST", "127.0.0.1")
PORT = int(os.environ.get("SKYKOI_KOI_BRAIN_PORT", "18794"))
TOKEN = os.environ.get("SKYKOI_KOI_BRAIN_TOKEN", "")

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from koibrain.config import load_config  # noqa: E402
from koibrain.orchestrator import Orchestrator  # noqa: E402

CONFIG = load_config()
# One orchestrator per session: conversation history must not bleed between a
# voice call, a Discord thread, and the TUI. The SQLite memory underneath is
# shared on purpose — cross-session recall is a feature, cross-session chatter
# is not (chat turns are never persisted, same rule as upstream).
SESSIONS: dict[str, Orchestrator] = {}
# One task at a time per session; different sessions run concurrently.
LOCKS: dict[str, asyncio.Lock] = {}


def _session(session_id: str) -> tuple[Orchestrator, asyncio.Lock]:
    if session_id not in SESSIONS:
        SESSIONS[session_id] = Orchestrator(CONFIG)
        LOCKS[session_id] = asyncio.Lock()
    return SESSIONS[session_id], LOCKS[session_id]


def _authorized(request: web.Request) -> bool:
    if not TOKEN:
        return True
    return request.headers.get("authorization", "") == f"Bearer {TOKEN}"


async def handle_run(request: web.Request) -> web.StreamResponse:
    if not _authorized(request):
        return web.json_response({"error": "unauthorized"}, status=401)
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "expected JSON"}, status=400)
    task = str(body.get("task") or "").strip()
    if not task:
        return web.json_response({"error": "expected { task }"}, status=400)
    session_id = str(body.get("sessionId") or "default")

    response = web.StreamResponse(
        headers={"Content-Type": "application/x-ndjson", "Cache-Control": "no-store"}
    )
    await response.prepare(request)

    async def emit(event: dict[str, Any]) -> None:
        try:
            await response.write((json.dumps(event) + "\n").encode())
        except Exception:
            pass  # caller hung up; the run still finishes and memory still records it

    orch, lock = _session(session_id)
    started = time.time()
    async with lock:
        orch.emit = emit
        try:
            result = await orch.handle(task)
            await emit({"kind": "done", "ok": True, "result": result, "ms": int((time.time() - started) * 1000)})
        except Exception as error:  # noqa: BLE001 — the caller gets the reason, not a hang
            traceback.print_exc()
            await emit({"kind": "done", "ok": False, "error": str(error), "ms": int((time.time() - started) * 1000)})
        finally:
            orch.emit = None
    await response.write_eof()
    return response


async def handle_health(request: web.Request) -> web.Response:
    return web.json_response(
        {
            "ok": True,
            "model": CONFIG.primary_model,
            "fallbacks": CONFIG.fallback_models,
            "policy": CONFIG.policy_mode,
            "sessions": len(SESSIONS),
        }
    )


async def _serve() -> None:
    app = web.Application(client_max_size=2 * 1024 * 1024)
    app.add_routes([web.post("/run", handle_run), web.get("/health", handle_health)])
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, HOST, PORT)
    await site.start()
    print(f"[koi-brain] listening on {HOST}:{PORT} · model {CONFIG.primary_model} · policy {CONFIG.policy_mode}", flush=True)
    try:
        await asyncio.Event().wait()
    finally:
        await runner.cleanup()


def main() -> None:
    # Proactor loop: subprocess spawning (run_shell tool) needs it on Windows,
    # same lesson as the browser agent.
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
    try:
        asyncio.run(_serve())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()