import { createHash, randomUUID } from "crypto";
import { Duplex } from "stream";
import { ConnectionUpgrader, WebSocketClient, WebSocketServerOptions, Request } from "./types.js";
import { Serializable } from "child_process";

export class WebSocketServer implements ConnectionUpgrader {
  public readonly clients: Map<string, WebSocketClient> = new Map();

  constructor(public readonly options: WebSocketServerOptions) {
    this.onUpgrade = this.onUpgrade.bind(this);
    this.isConnected = this.isConnected.bind(this);
    this.writeTo = this.writeTo.bind(this);
    this.broadcast = this.broadcast.bind(this);
  }

  public isConnected(clientUUID: string) {
    return this.clients.get(clientUUID) !== undefined;
  }

  public isUUIDValid(uuid: string) {
    return !this.clients.has(uuid);
  }

  public async broadcast(data: string, fromUUID?: string) {
    const tR = [];
    const clients = this.clients.values();
    for (const client of clients) {
      if (fromUUID === undefined || client.uuid !== fromUUID) {
        tR.push(this.writeTo(client.uuid, data));
      }
    }
    await Promise.allSettled(tR);
  }

  public async writeTo(clientUUID: string, data: string) {
    const client = this.clients.get(clientUUID);
    if (!client) {
      throw new Error("client not found");
    }
    if (!client.socket.writable) {
      this.clients.delete(clientUUID);
      throw new Error("socket not writable");
    }
    if (client.socket.writableLength > (this.options.maxFrameSize ?? DEFAULT_MAX_FRAME_SIZE)) {
      client.socket.destroy();
      this.clients.delete(clientUUID);
      throw new Error("client send buffer full");
    }
    return new Promise<void>((resolve, reject) => {
      try {
        client.socket.write(createFrame(data), (err?: Error | null) => {
          if (err) {
            try {
              client.socket.destroy(err);
            } catch (e) {
              console.error(e);
            }
            this.clients.delete(clientUUID);
            reject(err);
          } else {
            resolve();
          }
        });
      } catch (e) {
        try {
          client.socket.destroy();
        } catch (e) {
          console.error(e);
        }
        this.clients.delete(clientUUID);
        reject(e);
      }
    });
  }

  public async onUpgrade(req: Request, socket: Duplex, head: Buffer) {
    const acceptKey = req.headers["sec-websocket-key"];
    req.logger.debug("validating with acceptKey [%s]", acceptKey);
    if (
      !socket ||
      (!req.headers.connection || String(req.headers.connection).toLocaleLowerCase().indexOf("upgrade") === -1) ||
      String(req.headers.upgrade).toLocaleLowerCase() !== "websocket" ||
      String(req.method).toUpperCase() !== "GET" ||
      !acceptKey
    ) {
      req.logger.warn("bad request");
      socket.end("HTTP/1.1 400 Bad Request");
      socket.destroy();
      return;
    } else {
      const validateRet = this.options.validate ? await this.options.validate(req) : true;
      const uuid = typeof validateRet === "boolean" ?
        (validateRet ? randomUUID() : null) :
        (typeof validateRet === "object" ? validateRet.uuid :
          validateRet ? randomUUID() : null
        );
      const extraHeaders = typeof validateRet === "object" ? validateRet.headers : undefined;
      if (!uuid) {
        req.logger.debug("validating failed with result [%s]", validateRet);
        socket.end("HTTP/1.1 403 Forbidden");
        socket.destroy();
        return;
      } else {
        if (this.options.maxConnections !== undefined && this.clients.size >= this.options.maxConnections) {
          req.logger.warn("ignoring connection to many %s >= %s", this.clients.size, this.options.maxConnections);
          socket.end("HTTP/1.1 403 Forbidden");
          socket.destroy();
          return;
        }
        if (!this.isUUIDValid(uuid)) {
          req.logger.error("ignoring connection invalid uuid");
          socket.end("HTTP/1.1 403 Forbidden");
          socket.destroy();
          return;
        }
        const client: WebSocketClient = {
          uuid,
          head,
          req,
          socket
        };
        req.logger.debug("upgrading connection");

        socket.write(createUpgradeHeaders(acceptKey, extraHeaders));

        socket.on("data", async (data: Buffer) => {
          try {
            const frame = parseFrame(data, this.options.maxFrameSize);
            if (frame === PING) {
              try {
                socket.write(Buffer.from([0x8A, 0x00])); // pong frame
              } catch (e) {
                req.logger.error(e);
              }
            } else if (frame !== null) {
              if (!this.options.onMessage) {
                socket.end();
                socket.destroy();
              } else {
                try {
                  await this.options.onMessage(client, frame);
                } catch (e) {
                  req.logger.error(e);
                }
              }
            } else {
              // send close frame back and clean up
              try {
                socket.write(Buffer.from([0x88, 0x00]));
              } catch (e) {
                req.logger.error(e);
              }
              socket.end();
              socket.destroy();
            }
          } catch (e) {
            try {
              req.logger?.error(e);
              socket.end();
              socket.destroy();
            } catch (e) {
              console.error(e);
            }
          }
        });

        socket.on("error", async (error) => {
          req.logger.debug("upgraded connection error!");
          req.logger.error(error);
          if (this.options.onError) {
            try {
              await this.options.onError(client, error);
            } catch (e) {
              req.logger.error(e);
            }
          }
        });

        socket.on("end", async () => {
          req.logger.debug("upgraded connection disconnected!");
          this.clients.delete(uuid);
          if (this.options.onDisconnect) {
            try {
              await this.options.onDisconnect(client);
            } catch (e) {
              req.logger.error(e);
            }
          }
        });



        req.logger.debug("connection upgraded");
        this.clients.set(uuid, client);
        if (this.options.onConnection) {
          try {
            await this.options.onConnection(client);
          } catch (e) {
            req.logger.error(e);
          }
        }
        return;
      }
    }
  }
}

const ClusterWebSocketServerMessageType = "ClusterWebSocketServerMessage";

export interface ClusterWebSocketServerMessage {
  type: typeof ClusterWebSocketServerMessageType;
  action: "connection" | "disconnection" | "sendMessage" | "error";
  clientUUID?: string;
  errorMessage?: string;
  fromUUID?: string;
  payload?: Serializable;
}

export class ClusterWebSocketServer extends WebSocketServer {
  public static instance: ClusterWebSocketServer;
  public remoteClients: Set<string> = new Set();
  public constructor(options: WebSocketServerOptions) {
    super({
      ...options,
      validate: (req) => {
        if (
          this.options.maxConnections !== undefined &&
          this.clients.size + this.remoteClients.size >= this.options.maxConnections
        ) {
          return false;
        } else {
          return options.validate ? options.validate(req) : true;
        }
      },
      onError: async (req, error) => {
        if (process.send) {
          try {
            process.send({
              type: ClusterWebSocketServerMessageType,
              action: "error",
              clientUUID: req.uuid,
              errorMessage: error.message
            } as ClusterWebSocketServerMessage);
          } catch (e) {
            req.req.logger.error(e);
          }
        }
        if (options.onError) {
          await options.onError(req, error);
        }
      },
      onConnection: async (req) => {
        if (process.send) {
          try {
            process.send({
              type: ClusterWebSocketServerMessageType,
              action: "connection",
              clientUUID: req.uuid
            } as ClusterWebSocketServerMessage);
          } catch (e) {
            req.req.logger.error(e);
          }
        }
        if (options.onConnection) {
          await options.onConnection(req);
        }
      },
      onDisconnect: async (req) => {
        if (process.send) {
          try {
            process.send({
              type: ClusterWebSocketServerMessageType,
              action: "disconnection",
              clientUUID: req.uuid
            } as ClusterWebSocketServerMessage);
          } catch (e) {
            req.req.logger.error(e);
          }
        }
        if (options.onDisconnect) {
          await options.onDisconnect(req);
        }
      }
    });
    if (ClusterWebSocketServer.instance) {
      throw new Error("cannot create more than one instance");
    }
    ClusterWebSocketServer.instance = this;
    if (process.send) {
      process.on("message", async (data) => {
        try {
          const msg = (data as ClusterWebSocketServerMessage);
          if (
            msg &&
            msg.type === ClusterWebSocketServerMessageType &&
            msg.action) {
            // receive message from cluster workers
            switch (msg.action) {
              case "connection":
                if (!msg.clientUUID) {
                  throw new Error(`action [${msg.action}] without clientUUID`);
                }
                this.remoteClients.add(msg.clientUUID);
                break;
              case "disconnection":
                if (!msg.clientUUID) {
                  throw new Error(`action [${msg.action}] without clientUUID`);
                }
                this.remoteClients.delete(msg.clientUUID);
                break;
              case "sendMessage": {
                const payload = String(msg.payload);
                if (!msg.clientUUID) {
                  // broadcast to local clients
                  await super.broadcast(payload, msg.fromUUID);
                } else if (this.isConnected(msg.clientUUID)) {
                  // write if local client
                  await super.writeTo(msg.clientUUID, payload);
                }
                break;
              }
              default:
                throw new Error(`action [${msg.action}] not supported`);
            }
          }
        } catch (e) {
          console.error(e);
        }
      });
    }
  }

  public async broadcast(payload: string, fromUUID?: string): Promise<void> {
    if (process.send) {
      process.send({
        type: ClusterWebSocketServerMessageType,
        action: "sendMessage",
        payload,
        fromUUID
      } as ClusterWebSocketServerMessage);
    }
    return super.broadcast(payload, fromUUID);
  }

  public async writeTo(clientUUID: string, payload: string): Promise<void> {
    if (this.remoteClients.has(clientUUID)) {
      // write to remote client via IPC
      if (process.send) {
        process.send({
          type: ClusterWebSocketServerMessageType,
          action: "sendMessage",
          clientUUID,
          payload
        } as ClusterWebSocketServerMessage);
        return;
      }
    }
    // due to the implementation of isUUIDValid depends on this.remoteClients to be sync there can be two clients registered with the same UUID in the cluster
    if (this.clients.has(clientUUID)) {
      // write to local client
      return super.writeTo(clientUUID, payload);
    }
  }

  public isUUIDValid(uuid: string) {
    return !this.clients.has(uuid) && !this.remoteClients.has(uuid);
  }
}

/**
 * Private
 */

const PING = Symbol()
const OPCODES = { text: 0x01, close: 0x08, ping: 0x09, pong: 0x0A };
const GUID: string = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
//Set-Cookie: <cookie-name>=<cookie-value>; Domain=<domain-value>; Secure; HttpOnly
function createUpgradeHeaders(acceptKey: string, extraHeaders?: {
  name: string;
  value: string;
}[]) {
  const acceptValue = createHash('sha1')
    .update(acceptKey + GUID, 'binary')
    .digest('base64');
  let extra = extraHeaders && extraHeaders.length > 0 ? `\r\n${extraHeaders.map(h => `${h.name}: ${h.value}`).join("\r\n")}` : "";
  return `HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ${acceptValue}${extra}\r\n\r\n`;
}

const DEFAULT_MAX_FRAME_SIZE = 1024 * 1024; // 1 MiB

function parseFrame(buffer: Buffer, maxFrameSize = DEFAULT_MAX_FRAME_SIZE) {
  const firstByte = buffer.readUInt8(0);
  const opCode = firstByte & 0b00001111; // get last 4 bits of a byte

  if (opCode === OPCODES.close) {
    //this.emit('close');
    return null;
  } else if (opCode === OPCODES.ping) {
    return PING;
  } else if (opCode !== OPCODES.text) {
    return null;
  }

  const secondByte = buffer.readUInt8(1); // start with a payload length

  let offset = 2;
  let payloadLength = secondByte & 0b01111111; // get last 7 bits of a second byte

  if (payloadLength === 126) {
    payloadLength = buffer.readUInt16BE(offset);  // read real 16-bit length
    offset += 2;
  } else if (payloadLength === 127) {
    payloadLength = Number(buffer.readBigUInt64BE(offset));  // read real 64-bit length
    offset += 8;
  }

  if (payloadLength > maxFrameSize) {
    return null;
  }

  const isMasked = Boolean((secondByte >>> 7) & 0x1); // get first bit of a second byte

  if (isMasked) {
    const maskingKey = buffer.readUInt32BE(offset); // read 4-byte mask
    offset += 4;
    const payload = buffer.subarray(offset, offset + payloadLength);
    const result = unmask(payload, maskingKey);
    return result.toString('utf-8');
  }

  return buffer.subarray(offset, offset + payloadLength).toString('utf-8');
}

function unmask(payload: Buffer, maskingKey: number) {
  const result = Buffer.alloc(payload.byteLength);

  for (let i = 0; i < payload.byteLength; ++i) {
    const j = i % 4;
    const maskingKeyByteShift = j === 3 ? 0 : (3 - j) << 3;
    const maskingKeyByte = (maskingKeyByteShift === 0 ? maskingKey : maskingKey >>> maskingKeyByteShift) & 0b11111111;
    const transformedByte = maskingKeyByte ^ payload.readUInt8(i);
    result.writeUInt8(transformedByte, i);
  }

  return result;
}

function createFrame(payload: string) {

  const payloadByteLength = Buffer.byteLength(payload);
  let payloadBytesOffset = 2;
  let payloadLength = payloadByteLength;

  if (payloadByteLength > 65535) { // length value cannot fit in 2 bytes
    payloadBytesOffset += 8;
    payloadLength = 127;
  } else if (payloadByteLength > 125) {
    payloadBytesOffset += 2;
    payloadLength = 126;
  }

  const buffer = Buffer.alloc(payloadBytesOffset + payloadByteLength);

  // first byte
  buffer.writeUInt8(0b10000001, 0); // [FIN (1), RSV1 (0), RSV2 (0), RSV3 (0), Opode (0x01 - text frame)]

  buffer[1] = payloadLength; // second byte - actual payload size (if <= 125 bytes) or 126, or 127

  if (payloadLength === 126) { // write actual payload length as a 16-bit unsigned integer
    buffer.writeUInt16BE(payloadByteLength, 2);
  } else if (payloadLength === 127) { // write actual payload length as a 64-bit unsigned integer
    buffer.writeBigUInt64BE(BigInt(payloadByteLength), 2);
  }

  buffer.write(payload, payloadBytesOffset);
  return buffer;
}