const { WebSocketServer, WebSocket } = require("ws");
const { MAX_JSON_BYTES, MAX_AUDIO_PAYLOAD_BYTES, AUDIO_HEADER_BYTES, parseEnvelope, validateHello, parseAudioFrame, envelope } = require("./protocol");
const { insecureDeviceAllowed } = require("./device_store");

const MAX_CONTROL_MESSAGES_PER_SECOND = 120;
const MAX_AUDIO_MESSAGES_PER_SECOND = 200;
const MAX_SOURCE_MESSAGES_PER_SECOND = 1000;

class CompanionGateway {
  constructor(options) {
    this.devices = options.devices;
    this.sessions = options.sessions;
    this.log = options.log || { append() {} };
    this.wss = new WebSocketServer({ noServer: true, maxPayload: Math.max(MAX_JSON_BYTES, AUDIO_HEADER_BYTES + MAX_AUDIO_PAYLOAD_BYTES), perMessageDeflate: false, clientTracking: true });
    this.wss.on("connection", (socket, request, device) => this.connection(socket, request, device));
  }
  upgrade(request, socket, head) {
    const forwardedProto = String(request.headers["x-forwarded-proto"] || "").split(",")[0].trim();
    const proxyIsLocal = ["127.0.0.1", "::1", "::ffff:127.0.0.1"].includes(request.socket.remoteAddress);
    const secure = Boolean(request.socket.encrypted) || (proxyIsLocal && forwardedProto === "https");
    const auth = this.devices.authenticate(request.headers.authorization, "transcription.capture.v1");
    if (!auth.allowed) return reject(socket, auth.reason === "capability_revoked" ? 403 : 401, auth.reason);
    const requestOrigin = `http://${request.headers.host || "invalid"}`;
    if (!secure && !insecureDeviceAllowed(auth.device, requestOrigin, request.socket.remoteAddress)) return reject(socket, 426, "tls_required");
    const origin = request.headers.origin;
    if (origin && !sameHostOrigin(origin, request.headers.host)) return reject(socket, 403, "origin_rejected");
    this.wss.handleUpgrade(request, socket, head, (ws) => this.wss.emit("connection", ws, request, auth.device));
  }
  connection(socket, _request, device) {
    let session = null;
    let helloComplete = false;
    let lastPong = Date.now();
    let windowStarted = Date.now();
    let controlMessagesInWindow = 0;
    let audioMessagesInWindow = 0;
    let sourceMessagesInWindow = 0;
    let audioMessagesDropped = 0;
    let sourceMessagesDropped = 0;
    const sourceStates = new Map();
    let messageChain = Promise.resolve();
    const send = (type, payload, sessionId = session?.id || null) => {
      if (socket.readyState === WebSocket.OPEN) socket.send(JSON.stringify(envelope(type, payload, sessionId)));
    };
    const helloTimer = setTimeout(() => closeWith(socket, 4408, "hello_timeout"), 5000);
    const heartbeat = setInterval(() => {
      if (Date.now() - lastPong > 30000) return closeWith(socket, 4408, "heartbeat_timeout");
      send("status", { kind: "heartbeat", state: session?.state || "connecting" });
    }, 10000);
    helloTimer.unref?.(); heartbeat.unref?.();
    socket.on("message", (data, isBinary) => {
      const body = Buffer.from(data);
      messageChain = messageChain.then(async () => {
      try {
        if (Date.now() - windowStarted >= 1000) {
          if (audioMessagesDropped) this.log.append({ kind: "audio_rate_limited", device_id: device.id, session_id: session?.id, dropped_frames: audioMessagesDropped });
          if (sourceMessagesDropped) this.log.append({ kind: "source_rate_limited", device_id: device.id, session_id: session?.id, dropped_updates: sourceMessagesDropped });
          windowStarted = Date.now(); controlMessagesInWindow = 0; audioMessagesInWindow = 0; sourceMessagesInWindow = 0; audioMessagesDropped = 0; sourceMessagesDropped = 0;
        }
        if (isBinary) {
          if (!helloComplete || !session) throw coded("HELLO_REQUIRED", "Complete the handshake before sending audio.");
          audioMessagesInWindow += 1;
          if (audioMessagesInWindow > MAX_AUDIO_MESSAGES_PER_SECOND) { audioMessagesDropped += 1; return; }
          const result = await this.sessions.audio(session.id, parseAudioFrame(body));
          if (result.gap) send("metric", { kind: "sequence_gap", missing_frames: result.gap });
          return;
        }
        const message = parseEnvelope(body);
        if (helloComplete && message.type === "source_update") {
          const sourceId = String(message.payload?.source_uuid || "");
          const signature = JSON.stringify(message.payload || {});
          if (sourceId && sourceStates.get(sourceId) === signature) return;
          if (sourceId) {
            sourceStates.set(sourceId, signature);
            if (sourceStates.size > 2048) sourceStates.delete(sourceStates.keys().next().value);
          }
          sourceMessagesInWindow += 1;
          if (sourceMessagesInWindow > MAX_SOURCE_MESSAGES_PER_SECOND) { sourceMessagesDropped += 1; return; }
        } else {
          controlMessagesInWindow += 1;
          if (controlMessagesInWindow > MAX_CONTROL_MESSAGES_PER_SECOND) throw coded("RATE_LIMIT", "Companion control message rate exceeded its limit.");
        }
        if (!helloComplete) {
          const hello = validateHello(message);
          this.devices.updateRuntime(device.id, { companion_version: hello.companion_version, companion_plugin_version: hello.plugin_version });
          const created = this.sessions.create(device, send, hello.resume_session_id);
          session = created.session;
          helloComplete = true;
          clearTimeout(helloTimer);
          send("hello_ack", {
            protocol_version: 1, resumed: created.resumed,
            server_plugin_version: require("../../plugin.json").version,
            capabilities: ["transcription.server.v1", "settings.revision.v1", "pcm_s16le"],
            heartbeat_ms: 10000, recovery_window_ms: 5000
          });
          this.log.append({ kind: "connection", state: "authenticated", device_id: device.id, session_id: session.id });
          return;
        }
        await this.structured(session, message, send, () => { lastPong = Date.now(); });
      } catch (error) {
        send("error", { code: error.code || "INVALID_MESSAGE", message: error.message, recoverable: !["INCOMPATIBLE_VERSION", "HELLO_REQUIRED"].includes(error.code) });
        this.log.append({ kind: "protocol_error", device_id: device.id, session_id: session?.id, code: error.code || "INVALID_MESSAGE", message: error.message });
        if (["INCOMPATIBLE_VERSION", "HELLO_REQUIRED", "RATE_LIMIT"].includes(error.code)) closeWith(socket, 4400, error.code);
      }
      });
    });
    socket.on("close", () => {
      clearTimeout(helloTimer); clearInterval(heartbeat);
      if (session) this.sessions.disconnect(session.id);
      this.log.append({ kind: "connection", state: "closed", device_id: device.id, session_id: session?.id });
    });
    socket.on("error", (error) => this.log.append({ kind: "connection", state: "error", device_id: device.id, session_id: session?.id, message: error.message }));
  }
  async structured(session, message, send, pong) {
    switch (message.type) {
      case "ping":
        pong();
        send("pong", { received_id: message.id });
        send("status", { kind: "readiness", ...(await this.sessions.readiness()), server_plugin_version: require("../../plugin.json").version });
        break;
      case "source_update": {
        const source = this.sessions.updateSource(session.id, message.payload || {});
        this.devices.updateRuntime(session.deviceId, { bridge_installed: true, bridge_connected: true });
        send("status", { kind: "source", source });
        break;
      }
      case "obs_state": {
        const status = await this.sessions.updateObsState(session.id, message.payload || {});
        this.devices.updateRuntime(session.deviceId, status);
        send("status", { kind: "obs", ...status });
        break;
      }
      case "readiness": send("status", { kind: "readiness", ...(await this.sessions.readiness()), server_plugin_version: require("../../plugin.json").version }); break;
      case "start": send("status", { kind: "session", ...(await this.sessions.start(session.id, message.payload || {})) }); break;
      case "stop": send("status", { kind: "session", ...(await this.sessions.stop(session.id, cleanReason(message.payload?.reason))) }); break;
      case "ack": break;
      default: throw coded("UNEXPECTED_MESSAGE", `Message ${message.type} is not valid after the handshake.`);
    }
  }
  async close() {
    for (const client of this.wss.clients) closeWith(client, 1001, "plugin_shutdown");
    await new Promise((resolve) => this.wss.close(resolve));
  }
}

function reject(socket, status, reason) { const labels = { 401: "Unauthorized", 403: "Forbidden", 426: "Upgrade Required" }; socket.write(`HTTP/1.1 ${status} ${labels[status] || "Rejected"}\r\nConnection: close\r\nContent-Type: application/json\r\n\r\n${JSON.stringify({ error: reason })}`); socket.destroy(); }
function closeWith(socket, code, reason) { if (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING) socket.close(code, String(reason).slice(0, 120)); }
function sameHostOrigin(origin, host) { try { return new URL(origin).host === host; } catch { return false; } }
function coded(code, message) { return Object.assign(new Error(message), { code }); }
function cleanReason(value) { return ["requested", "test_complete", "benchmark_complete", "silence_timeout", "disconnect"].includes(String(value)) ? String(value) : "requested"; }

module.exports = { CompanionGateway, sameHostOrigin, MAX_CONTROL_MESSAGES_PER_SECOND, MAX_AUDIO_MESSAGES_PER_SECOND, MAX_SOURCE_MESSAGES_PER_SECOND };
