import { createReadStream, existsSync } from "node:fs";
import { stat } from "node:fs/promises";
import { createServer, type IncomingMessage, type ServerResponse } from "node:http";
import path from "node:path";
import { WebSocket, WebSocketServer } from "ws";

type TransportLogEntry =
  | {
      type: "connection-open";
      timestamp: number;
      connectionId: number;
    }
  | {
      type: "connection-close";
      timestamp: number;
      connectionId: number;
      code: number;
      reason: string;
    }
  | {
      type: "binary-frame";
      timestamp: number;
      connectionId: number;
      byteLength: number;
    }
  | {
      type: "text-frame";
      timestamp: number;
      connectionId: number;
      text: string;
    }
  | {
      type: "finish";
      timestamp: number;
      connectionId: number;
    }
  | {
      type: "server-message";
      timestamp: number;
      payload: Record<string, unknown>;
    }
  | {
      type: "outbound-frame-delayed";
      timestamp: number;
      connectionId: number;
      frameNumber: number;
      delayMs: number;
      kind: "server-message" | "finish";
    }
  | {
      type: "outbound-frame-dropped";
      timestamp: number;
      connectionId: number;
      frameNumber: number;
      kind: "server-message" | "finish";
    }
  | {
      type: "drain-stalled";
      timestamp: number;
      durationMs: number;
    };
type OutboundFrame = {
  socket: WebSocket;
  connectionId: number;
  payload: string;
  kind: "server-message" | "finish";
  deliverAt: number;
  frameNumber: number;
};

type DrainFaultState = {
  delayedFramesMs: number;
  dropEveryNthFrame: number;
  drainStalledUntilMs: number;
};

const CONTENT_TYPES: Record<string, string> = {
  ".html": "text/html; charset=utf-8",
  ".js": "application/javascript; charset=utf-8",
  ".map": "application/json; charset=utf-8",
  ".css": "text/css; charset=utf-8",
};

const waitFor = async <T>(
  getter: () => T | undefined,
  timeoutMs = 10_000,
  errorMessage = "Timed out waiting for condition."
): Promise<T> => {
  const startedAt = Date.now();

  while (Date.now() - startedAt < timeoutMs) {
    const value = getter();
    if (value !== undefined) {
      return value;
    }

    await new Promise((resolve) => setTimeout(resolve, 50));
  }

  throw new Error(errorMessage);
};

export class ScriptedTranscriptionServer {
  private readonly sockets = new Map<number, WebSocket>();
  private readonly transportLog: TransportLogEntry[] = [];
  private readonly pendingOutboundFrames: OutboundFrame[] = [];
  private httpServer;
  private wsServer;
  private nextConnectionId = 1;
  private port = 0;
  private outboundFrameCount = 0;
  private outboundFlushTimer: ReturnType<typeof setTimeout> | null = null;
  private faultState: DrainFaultState = {
    delayedFramesMs: 0,
    dropEveryNthFrame: 0,
    drainStalledUntilMs: 0,
  };

  constructor(private readonly rootDir: string) {
    this.httpServer = createServer(this.handleRequest);
    this.wsServer = new WebSocketServer({ noServer: true });
  }

  public async start() {
    this.httpServer.on("upgrade", (request, socket, head) => {
      if (!request.url?.startsWith("/ws")) {
        socket.destroy();
        return;
      }

      this.wsServer.handleUpgrade(request, socket, head, (websocket) => {
        this.wsServer.emit("connection", websocket, request);
      });
    });

    this.wsServer.on("connection", (socket) => {
      const connectionId = this.nextConnectionId++;
      this.sockets.set(connectionId, socket);
      this.transportLog.push({
        type: "connection-open",
        timestamp: Date.now(),
        connectionId,
      });

      socket.on("message", (data, isBinary) => {
        if (isBinary) {
          this.transportLog.push({
            type: "binary-frame",
            timestamp: Date.now(),
            connectionId,
            byteLength: data.byteLength,
          });
          return;
        }

        const text = data.toString();
        this.transportLog.push({
          type: "text-frame",
          timestamp: Date.now(),
          connectionId,
          text,
        });

        try {
          const payload = JSON.parse(text);
          if (payload?.action === "finish") {
            this.transportLog.push({
              type: "finish",
              timestamp: Date.now(),
              connectionId,
            });
            this.queueOutboundFrame({
              socket,
              connectionId,
              payload: JSON.stringify({ is_partial: false, data: {} }),
              kind: "finish",
            });
          }
        } catch (error) {
          // Ignore malformed control text in the scaffold server.
        }
      });

      socket.on("close", (code, reason) => {
        this.transportLog.push({
          type: "connection-close",
          timestamp: Date.now(),
          connectionId,
          code,
          reason: reason.toString(),
        });
        this.sockets.delete(connectionId);
      });
    });

    await new Promise<void>((resolve, reject) => {
      this.httpServer.once("error", reject);
      this.httpServer.listen(0, "127.0.0.1", () => {
        const address = this.httpServer.address();
        if (!address || typeof address === "string") {
          reject(new Error("Failed to resolve the scripted test server port."));
          return;
        }

        this.port = address.port;
        resolve();
      });
    });
  }

  public async stop() {
    for (const socket of this.sockets.values()) {
      socket.terminate();
    }
    this.sockets.clear();

    this.httpServer.closeIdleConnections?.();
    this.httpServer.closeAllConnections?.();

    await new Promise<void>((resolve) => this.wsServer.close(() => resolve()));
    await new Promise<void>((resolve, reject) =>
      this.httpServer.close((error) => (error ? reject(error) : resolve()))
    );
  }

  public getHttpBaseUrl() {
    return `http://127.0.0.1:${this.port}`;
  }

  public getWsUrl() {
    return `ws://127.0.0.1:${this.port}/ws`;
  }

  public getHarnessUrl() {
    return `${this.getHttpBaseUrl()}/e2e/pages/resilience-harness.html?ws=${encodeURIComponent(
      this.getWsUrl()
    )}`;
  }

  public getTransportLog() {
    return this.transportLog.slice();
  }

  public clearTransportLog() {
    this.transportLog.length = 0;
  }

  public clearFaults() {
    this.faultState = {
      delayedFramesMs: 0,
      dropEveryNthFrame: 0,
      drainStalledUntilMs: 0,
    };
    this.pendingOutboundFrames.length = 0;
    this.outboundFrameCount = 0;
    if (this.outboundFlushTimer) {
      clearTimeout(this.outboundFlushTimer);
      this.outboundFlushTimer = null;
    }
  }

  public async waitForConnectionCount(count: number, timeoutMs = 10_000) {
    return waitFor(
      () => {
        const opened = this.transportLog.filter(
          (entry) => entry.type === "connection-open"
        );
        return opened.length >= count ? opened.length : undefined;
      },
      timeoutMs,
      `Timed out waiting for ${count} websocket connection(s).`
    );
  }

  public async waitForBinaryFrames(count: number, timeoutMs = 10_000) {
    return waitFor(
      () => {
        const frames = this.transportLog.filter(
          (entry) => entry.type === "binary-frame"
        );
        return frames.length >= count ? frames.length : undefined;
      },
      timeoutMs,
      `Timed out waiting for ${count} binary audio frame(s).`
    );
  }

  public async waitForFinish(timeoutMs = 10_000) {
    return waitFor(
      () =>
        this.transportLog.find((entry) => entry.type === "finish") as
          | Extract<TransportLogEntry, { type: "finish" }>
          | undefined,
      timeoutMs,
      "Timed out waiting for the finish control message."
    );
  }

  public sendPartial(text: string) {
    this.broadcast({
      is_partial: true,
      data: { text },
    });
  }

  public sendFinal(text: string) {
    this.broadcast({
      is_partial: false,
      data: { text },
    });
  }

  public delayFrames(milliseconds: number) {
    this.faultState.delayedFramesMs = Math.max(0, milliseconds);
  }

  public dropEveryNthFrame(frameCount: number) {
    this.faultState.dropEveryNthFrame = Math.max(0, Math.floor(frameCount));
  }

  public stallDrain(durationMs: number) {
    const untilMs = Date.now() + Math.max(0, durationMs);
    this.faultState.drainStalledUntilMs = Math.max(
      this.faultState.drainStalledUntilMs,
      untilMs
    );
    this.transportLog.push({
      type: "drain-stalled",
      timestamp: Date.now(),
      durationMs,
    });
    this.scheduleOutboundFlush();
  }

  public terminateActiveSockets() {
    for (const socket of this.sockets.values()) {
      socket.terminate();
    }
  }

  private broadcast(payload: Record<string, unknown>) {
    this.transportLog.push({
      type: "server-message",
      timestamp: Date.now(),
      payload,
    });

    const encoded = JSON.stringify(payload);
    for (const socket of this.sockets.values()) {
      if (socket.readyState === WebSocket.OPEN) {
        this.queueOutboundFrame({
          socket,
          connectionId: this.getConnectionId(socket),
          payload: encoded,
          kind: "server-message",
        });
      }
    }
  }

  private queueOutboundFrame(frame: Omit<OutboundFrame, "deliverAt" | "frameNumber">) {
    const frameNumber = ++this.outboundFrameCount;

    if (
      frame.kind === "server-message" &&
      this.faultState.dropEveryNthFrame > 0 &&
      frameNumber % this.faultState.dropEveryNthFrame === 0
    ) {
      this.transportLog.push({
        type: "outbound-frame-dropped",
        timestamp: Date.now(),
        connectionId: frame.connectionId,
        frameNumber,
        kind: frame.kind,
      });
      return;
    }

    const now = Date.now();
    const deliverAt = Math.max(
      now + this.faultState.delayedFramesMs,
      this.faultState.drainStalledUntilMs
    );

    if (deliverAt > now) {
      this.transportLog.push({
        type: "outbound-frame-delayed",
        timestamp: now,
        connectionId: frame.connectionId,
        frameNumber,
        delayMs: deliverAt - now,
        kind: frame.kind,
      });
    }

    if (deliverAt <= now && this.pendingOutboundFrames.length === 0) {
      this.sendOutboundFrame({
        ...frame,
        deliverAt: now,
        frameNumber,
      });
      return;
    }

    this.pendingOutboundFrames.push({
      ...frame,
      deliverAt,
      frameNumber,
    });
    this.scheduleOutboundFlush();
  }

  private scheduleOutboundFlush() {
    if (this.outboundFlushTimer || this.pendingOutboundFrames.length === 0) {
      return;
    }

    const nextDeliverAt = Math.min(
      ...this.pendingOutboundFrames.map((frame) => frame.deliverAt)
    );
    const delayMs = Math.max(0, nextDeliverAt - Date.now());

    this.outboundFlushTimer = setTimeout(() => {
      this.outboundFlushTimer = null;
      void this.flushOutboundFrames();
    }, delayMs);
  }

  private async flushOutboundFrames() {
    const now = Date.now();
    const readyFrames = this.pendingOutboundFrames.filter(
      (frame) => frame.deliverAt <= now
    );
    this.pendingOutboundFrames.splice(
      0,
      this.pendingOutboundFrames.length,
      ...this.pendingOutboundFrames.filter((frame) => frame.deliverAt > now)
    );

    for (const frame of readyFrames) {
      this.sendOutboundFrame(frame);
    }

    if (this.pendingOutboundFrames.length > 0) {
      this.scheduleOutboundFlush();
    }
  }

  private sendOutboundFrame(frame: OutboundFrame) {
    if (frame.socket.readyState !== WebSocket.OPEN) {
      return;
    }

    frame.socket.send(frame.payload);
  }

  private getConnectionId(socket: WebSocket) {
    for (const [connectionId, currentSocket] of this.sockets.entries()) {
      if (currentSocket === socket) {
        return connectionId;
      }
    }

    return 0;
  }

  private handleRequest = async (
    request: IncomingMessage,
    response: ServerResponse
  ) => {
    try {
      const requestUrl = new URL(
        request.url || "/",
        "http://127.0.0.1"
      );
      const relativePath =
        requestUrl.pathname === "/"
          ? "/e2e/pages/resilience-harness.html"
          : requestUrl.pathname;
      const resolvedPath = path.resolve(this.rootDir, `.${relativePath}`);

      if (!resolvedPath.startsWith(this.rootDir) || !existsSync(resolvedPath)) {
        response.writeHead(404);
        response.end("Not found");
        return;
      }

      const fileStat = await stat(resolvedPath);
      if (!fileStat.isFile()) {
        response.writeHead(404);
        response.end("Not found");
        return;
      }

      const extension = path.extname(resolvedPath);
      response.writeHead(200, {
        "Content-Type": CONTENT_TYPES[extension] || "application/octet-stream",
      });
      createReadStream(resolvedPath).pipe(response);
    } catch (error) {
      response.writeHead(500);
      response.end(
        `Scripted transcription server failed to serve ${request.url}: ${
          error instanceof Error ? error.message : String(error)
        }`
      );
    }
  };
}