import { WebSocketServer } from "ws";
import { randomUUID } from "node:crypto";
import {
  PROTOCOL_VERSION,
  safeParseClientMessageJson,
  type ServerMessage,
} from "@agent-client/shared";
import { runAgent } from "./agent.js";

const PORT = Number(process.env.WS_PORT ?? 4291);
const SERVICE_TOKEN = process.env.GATEWAY_TOKEN ?? "";
const PROJECT_DIR = process.env.AGENT_CWD ?? process.cwd();

if (!SERVICE_TOKEN) {
  console.error("GATEWAY_TOKEN is required (used for WS auth too in v0.1).");
  process.exit(1);
}

const wss = new WebSocketServer({ port: PORT });

console.log(`[service] listening ws://localhost:${PORT}`);
console.log(`[service] AGENT_CWD = ${PROJECT_DIR}`);
console.log(
  `[service] ANTHROPIC_BASE_URL = ${process.env.ANTHROPIC_BASE_URL ?? "<unset>"}`,
);

wss.on("connection", (ws, req) => {
  // Token comes in via ?token=... since browsers can't set headers on ws.
  const url = new URL(req.url ?? "/", `http://${req.headers.host}`);
  if (url.searchParams.get("token") !== SERVICE_TOKEN) {
    ws.close(1008, "unauthorized");
    return;
  }

  const connId = randomUUID().slice(0, 8);
  console.log(`[service] +conn ${connId}`);

  const send = (msg: ServerMessage) => {
    if (ws.readyState === ws.OPEN) {
      ws.send(JSON.stringify(msg));
    }
  };

  send({ type: "ready", protocol: PROTOCOL_VERSION });

  // One agent run per active turn. Multiple turns are sequential per conn.
  let activeAbort: AbortController | null = null;

  ws.on("message", async (raw) => {
    const parsed = safeParseClientMessageJson(raw.toString());
    if (!parsed.success) {
      send({ type: "error", message: `invalid message: ${parsed.reason}` });
      return;
    }
    const msg = parsed.data;

    if (msg.type === "abort") {
      activeAbort?.abort();
      activeAbort = null;
      return;
    }

    if (activeAbort) {
      send({
        type: "error",
        id: msg.id,
        message: "上一轮还在执行中，请先取消",
      });
      return;
    }

    const ctrl = new AbortController();
    activeAbort = ctrl;
    try {
      await runAgent({
        id: msg.id,
        prompt: msg.prompt,
        cwd: PROJECT_DIR,
        abortController: ctrl,
        onMessage: send,
      });
    } finally {
      activeAbort = null;
    }
  });

  ws.on("close", () => {
    activeAbort?.abort();
    activeAbort = null;
    console.log(`[service] -conn ${connId}`);
  });
});

const shutdown = () => {
  console.log("[service] shutting down");
  wss.close(() => process.exit(0));
};
process.on("SIGINT", shutdown);
process.on("SIGTERM", shutdown);