import { WebSocketServer, WebSocket } from "ws";
import { SessionManager } from "./SessionManager";

declare global {
  var wss: WebSocketServer | undefined;
}

export function initWebSocketServer(port = 3001) {
  // Prevent duplicate WebSocket servers during hot reloads in development
  if (global.wss) {
    console.log("[WebSocketServer] WebSocket Server already active on port 3001 (cached).");
    return;
  }

  console.log(`[WebSocketServer] Launching NeuralLoom WebSocket Gateway on port ${port}...`);

  try {
    const wss = new WebSocketServer({ port });
    global.wss = wss;

    wss.on("connection", (ws: WebSocket & { sessionId?: string }) => {
      console.log("[WebSocketServer] New client connected via gateway.");

      let activeSessionId: string | null = null;
      let outputListener: ((data: string) => void) | null = null;
      let mockBuffer = "";

      ws.on("message", async (rawMsg) => {
        try {
          const msg = JSON.parse(rawMsg.toString());

          // Client handshake to establish connection to a specific session
          if (msg.type === "handshake") {
            const { sessionId } = msg;
            const manager = SessionManager.getInstance();
            const session = manager.getSession(sessionId);

            if (session) {
              activeSessionId = sessionId;
              ws.sessionId = sessionId;
              console.log(`[WebSocketServer] Client connected to active session: ${sessionId}`);

              // Subscribe to runner output events
              outputListener = (data: string) => {
                if (ws.readyState === WebSocket.OPEN) {
                  ws.send(JSON.stringify({ type: "output", data }));
                }
              };
              session.onOutput(outputListener);

              // Send handshake acknowledgment and full historical logs
              ws.send(
                JSON.stringify({
                  type: "handshake-ack",
                  sessionId,
                  history: session.logs.join(""),
                })
              );
            } else {
              ws.send(JSON.stringify({ type: "error", message: "Requested session not found" }));
            }
          }

          // User terminal input routing
          if (msg.type === "input" && activeSessionId) {
            const manager = SessionManager.getInstance();
            const session = manager.getSession(activeSessionId);
            if (session && session.status === "running") {
              if (session.type === "mock") {
                // Buffer input for simulated command parsing in MockRunner
                const data = msg.data;
                if (data === "\r" || data === "\r\n") {
                  ws.send(JSON.stringify({ type: "output", data: "\r\n" }));
                  const cmd = mockBuffer;
                  mockBuffer = "";
                  await session.sendInput(cmd);
                } else if (data === "\u007f" || data === "\b") { // Backspace
                  if (mockBuffer.length > 0) {
                    mockBuffer = mockBuffer.slice(0, -1);
                    ws.send(JSON.stringify({ type: "output", data: "\b \b" }));
                  }
                } else {
                  mockBuffer += data;
                  ws.send(JSON.stringify({ type: "output", data })); // Local echo of key
                }
              } else {
                // Pass raw keystrokes straight to Claude Code PTY
                await session.sendInput(msg.data, true);
              }
            }
          }
        } catch (err) {
          console.error("[WebSocketServer] Error processing socket packet:", err);
        }
      });

      ws.on("close", () => {
        if (activeSessionId && outputListener) {
          const manager = SessionManager.getInstance();
          const session = manager.getSession(activeSessionId);
          if (session) {
            session.offOutput(outputListener);
          }
        }
        console.log("[WebSocketServer] Client connection closed.");
      });

      ws.on("error", (err) => {
        console.error("[WebSocketServer] Socket error:", err);
      });
    });

    console.log(`[WebSocketServer] WebSocket gateway successfully listening on port ${port}.`);
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : "Unknown error";
    console.error(`[WebSocketServer] Failed to initialize WebSocket Server: ${errorMessage}`);
  }
}
