import {
  ArrayBuffer,
  ArrayBufferPrototypeGetResizable as resizable,
  ArrayPrototypePush as push,
  getter,
  ObjectCreate,
  ObjectDefineProperty as define,
  ReflectApply as apply,
  ReflectConstruct as construct,
  RegExp,
  RegExpPrototypeExec as exec,
  Set,
  SetPrototypeAdd as add,
  SetPrototypeHas as has,
  StringPrototypeSlice as slice,
  SymbolIterator,
  TypedArrayPrototypeGetBuffer as buffer,
  TypedArrayPrototypeGetLength as length,
  TypedArrayPrototypeGetTag as tag,
  TypeError,
  uncurryThis,
  WeakMap,
  WeakMapPrototypeSet as set,
} from "ext:core/primordials";
import {
  markPromiseHandled,
  promiseThen as then,
  safePromise,
} from "ext:core/promises";
import {
  dictionary,
  isObject,
  required,
  state,
  string,
  usv,
} from "ext:web/idl.js";
import { blobBody, fromBytes, snapshot } from "ext:blob/blob.js";
import { MessageEvent } from "ext:web/messaging.js";
import { userAgent } from "ext:web/navigator.js";
import {
  addListener,
  createException,
  dispatch,
  ErrorEvent,
  Event,
  EventTarget,
  preventDefault,
  removeListener,
  reportErrorValue,
} from "ext:web/events.js";
import { TextEncoder } from "ext:encoding/main.js";
import { URL } from "ext:url/main.js";
import {
  op_ws_close,
  op_ws_create,
  op_ws_fail,
  op_ws_retire,
  op_ws_send,
  op_ws_wait,
} from "ext:core/ops";

import type { WebSocketEvent } from "ext:core/ops";

import { byteLength as sharedLength } from "ext:core/shared_array_buffer";

const CONNECTING = 0,
  OPEN = 1,
  CLOSING = 2,
  CLOSED = 3;
type ReadyState =
  | typeof CONNECTING
  | typeof OPEN
  | typeof CLOSING
  | typeof CLOSED;
type EventType = "open" | "message" | "error" | "close";
type EventHandler<E extends Event> = (this: WebSocket, event: E) => unknown;
interface Handler {
  value: Function | null;
  listener: ((event: Event) => void) | null;
}
interface SocketState {
  id: number;
  href: string;
  origin: string;
  state: ReadyState;
  protocol: string;
  extensions: string;
  binaryType: "blob" | "arraybuffer";
  buffered: number;
  handlers: Record<EventType, Handler>;
}
interface CloseState {
  wasClean: boolean;
  code: number;
  reason: string;
}
const sockets = new WeakMap<object, SocketState>(),
  closeEvents = new WeakMap<object, CloseState>();
const encoder = new TextEncoder(),
  encode = uncurryThis(TextEncoder.prototype.encode);
const invalidToken = RegExp("[^A-Za-z0-9!#$%&'*+\\-.^_`|~]");
const floor = Math.floor,
  includes = uncurryThis(String.prototype.includes),
  urlHref = getter(URL.prototype, "href"),
  urlProtocol = getter(URL.prototype, "protocol"),
  urlOrigin = getter(URL.prototype, "origin");
// As in blob.js: the byteLength getter throws for anything but an ArrayBuffer.
const isView = ArrayBuffer.isView,
  byteLength = getter(ArrayBuffer.prototype, "byteLength"),
  dataViewBuffer = getter(DataView.prototype, "buffer");

function isBufferSource(value: unknown): boolean {
  if (isView(value)) return true;
  try {
    byteLength(value);
    return true;
  } catch {}
  try {
    sharedLength(value);
    return true;
  } catch {
    return false;
  }
}

function syntaxError(message: string): Error {
  return createException({ message, name: "SyntaxError" });
}

function toUint16(value: unknown): number {
  return (+(value as number) >>> 0) % 65536;
}

function clamp(value: unknown): number {
  const number = +(value as number);
  if (number !== number || number <= 0) return 0;
  if (number >= 65535) return 65535;
  const base = floor(number),
    fraction = number - base;
  return fraction < 0.5 || (fraction === 0.5 && base % 2 === 0)
    ? base
    : base + 1;
}

function convertProtocols(value: unknown): string[] {
  const list: string[] = [];
  if (value === undefined) return list;
  const method = isObject(value) ? value[SymbolIterator] : undefined;
  if (method === undefined || method === null) {
    push(list, string(value));
  } else {
    if (typeof method !== "function") {
      throw new TypeError("Invalid protocol iterator");
    }
    const iterator = apply(method, value, []);
    if (!isObject(iterator)) throw new TypeError("Invalid protocol iterator");
    const next = iterator.next;
    if (typeof next !== "function") {
      throw new TypeError("Invalid protocol iterator");
    }
    while (true) {
      const step = apply(next, iterator, []);
      if (!isObject(step)) {
        throw new TypeError("Invalid protocol iterator result");
      }
      if (step.done) break;
      push(list, string(step.value));
    }
  }
  return list;
}

// The pump holds the socket strongly until the close event, which is what keeps
// an OPEN socket with listeners from being collected.
async function pump(id: number, socket: WebSocket): Promise<void> {
  let failure: string | undefined;
  try {
    while (true) {
      let event: WebSocketEvent;
      try {
        event = await safePromise(op_ws_wait(id));
      } catch (error) {
        failure = usv(error);
        event = { kind: "close", clean: false, code: 1006, reason: "" };
      }
      const s = state(sockets, socket);
      try {
        if (event.kind === "open") {
          // close() may have been called while the connection was still establishing.
          if (s.state === CONNECTING) {
            s.state = OPEN;
            if (event.protocol !== "") s.protocol = event.protocol;
            dispatch(socket, new Event("open"));
          }
        } else if (event.kind === "text" || event.kind === "binary") {
          if (s.state === OPEN) {
            const data = event.kind === "text"
              ? event.data
              : s.binaryType === "blob"
              ? fromBytes(event.data, "")
              : buffer(event.data);
            dispatch(
              socket,
              new MessageEvent("message", { data, origin: s.origin }),
            );
          }
        } else if (event.kind === "closing") {
          if (s.state === OPEN) s.state = CLOSING;
        } else if (event.kind === "error") {
          failure = event.message;
        } else {
          s.state = CLOSED;
          if (failure !== undefined) {
            dispatch(socket, new ErrorEvent("error", { message: failure }));
          }
          dispatch(
            socket,
            new CloseEvent("close", {
              wasClean: event.clean,
              code: event.code,
              reason: event.reason,
            }),
          );
          return;
        }
      } catch (error) {
        reportErrorValue(error);
        if (s.state === CLOSED) return;
      }
    }
  } finally {
    op_ws_retire(id);
  }
}

const eventTypes: EventType[] = ["open", "message", "error", "close"];

function initializeSocket(
  socket: WebSocket,
  id: number,
  record: URL,
): WebSocket {
  const handlers = ObjectCreate(null) as Record<EventType, Handler>;
  for (let i = 0; i < eventTypes.length; i++) {
    handlers[eventTypes[i]] = { value: null, listener: null };
  }
  set(sockets, socket, {
    id,
    href: urlHref(record),
    origin: urlOrigin(record),
    state: CONNECTING,
    protocol: "",
    extensions: "",
    binaryType: "blob",
    buffered: 0,
    handlers,
  });
  markPromiseHandled(pump(id, socket));
  return socket;
}

function adoptWebSocket(id: number, href: string): WebSocket {
  const socket = construct(EventTarget, [], WebSocket) as WebSocket;
  return initializeSocket(socket, id, new URL(`ws${slice(href, 4)}`));
}

function handler(type: EventType) {
  return {
    get(this: WebSocket): Function | null {
      return state(sockets, this).handlers[type].value;
    },
    set(this: WebSocket, rawValue: unknown): void {
      const handler = state(sockets, this).handlers[type];
      const value = typeof rawValue === "function" ? rawValue : null;
      handler.value = value;
      if (value && !handler.listener) {
        handler.listener = (event) => {
          if (apply(handler.value!, this, [event]) === false) {
            preventDefault(event);
          }
        };
        addListener(this, type, handler.listener);
      } else if (!value && handler.listener) {
        removeListener(this, type, handler.listener);
        handler.listener = null;
      }
    },
    configurable: true,
  };
}

class CloseEvent extends Event {
  // @ts-expect-error Reflect.construct preserves new.target without a guest-replaceable super() call.
  constructor(type: unknown, init: unknown = {}) {
    required(arguments.length, 1);
    type = string(type);
    const options = dictionary(init);
    const self = construct(Event, [type, options], new.target) as CloseEvent;
    const rawCode = options.code,
      reason = options.reason;
    set(closeEvents, self, {
      wasClean: !!options.wasClean,
      code: rawCode === undefined ? 0 : toUint16(rawCode),
      reason: reason === undefined ? "" : usv(reason),
    });
    return self;
  }
  get wasClean() {
    return state(closeEvents, this).wasClean;
  }
  get code() {
    return state(closeEvents, this).code;
  }
  get reason() {
    return state(closeEvents, this).reason;
  }
}

interface WebSocket {
  readonly CONNECTING: typeof CONNECTING;
  readonly OPEN: typeof OPEN;
  readonly CLOSING: typeof CLOSING;
  readonly CLOSED: typeof CLOSED;
  onopen: EventHandler<Event> | null;
  onmessage: EventHandler<MessageEvent> | null;
  onerror: EventHandler<Event> | null;
  onclose: EventHandler<CloseEvent> | null;
}
class WebSocket extends EventTarget {
  declare static readonly CONNECTING: typeof CONNECTING;
  declare static readonly OPEN: typeof OPEN;
  declare static readonly CLOSING: typeof CLOSING;
  declare static readonly CLOSED: typeof CLOSED;

  // @ts-expect-error Reflect.construct preserves new.target without a guest-replaceable super() call.
  constructor(url: unknown, protocols: unknown = undefined) {
    required(arguments.length, 1);
    const self = construct(EventTarget, [], new.target) as WebSocket;
    url = usv(url);
    const list = convertProtocols(protocols);
    let record;
    try {
      record = new URL(url);
    } catch {
      throw syntaxError("The URL is invalid");
    }
    const protocol = urlProtocol(record);
    if (protocol === "http:") {
      record = new URL(`ws${slice(urlHref(record), 4)}`);
    } else if (protocol === "https:") {
      record = new URL(`wss${slice(urlHref(record), 5)}`);
    } else if (protocol !== "ws:" && protocol !== "wss:") {
      throw syntaxError("The URL's scheme must be ws, wss, http, or https");
    }
    const href = urlHref(record);
    if (includes(href, "#")) {
      throw syntaxError("The URL contains a fragment");
    }
    const seen = new Set<string>();
    for (let i = 0; i < list.length; i++) {
      const item = list[i];
      if (item === "" || exec(invalidToken, item) !== null) {
        throw syntaxError("Invalid WebSocket subprotocol");
      }
      if (has(seen, item)) {
        throw syntaxError("Duplicate WebSocket subprotocol");
      }
      add(seen, item);
    }
    // Permission failures surface synchronously, as in fetch.
    const id = op_ws_create([href, list, userAgent]);
    return initializeSocket(self, id, record);
  }
  get url() {
    return state(sockets, this).href;
  }
  get readyState() {
    return state(sockets, this).state;
  }
  get bufferedAmount() {
    return state(sockets, this).buffered;
  }
  get extensions() {
    return state(sockets, this).extensions;
  }
  get protocol() {
    return state(sockets, this).protocol;
  }
  get binaryType(): "blob" | "arraybuffer" {
    return state(sockets, this).binaryType;
  }
  set binaryType(value: unknown) {
    const s = state(sockets, this), binaryType = string(value);
    if (binaryType === "blob" || binaryType === "arraybuffer") {
      s.binaryType = binaryType;
    }
  }
  close(code: unknown = undefined, reason: unknown = undefined): void {
    const s = state(sockets, this);
    let closeCode: number | null = null;
    if (code !== undefined) closeCode = clamp(code);
    if (reason !== undefined) reason = usv(reason);
    if (
      closeCode !== null && closeCode !== 1000 &&
      (closeCode < 3000 || closeCode > 4999)
    ) {
      throw createException({
        message: "Invalid WebSocket close code",
        name: "InvalidAccessError",
      });
    }
    if (reason !== undefined && length(encode(encoder, reason)) > 123) {
      throw syntaxError("The WebSocket close reason is too long");
    }
    if (s.state === CLOSING || s.state === CLOSED) return;
    if (s.state === CONNECTING) {
      // Failing the connection instead of starting the closing handshake.
      s.state = CLOSING;
      op_ws_fail(s.id);
      return;
    }
    s.state = CLOSING;
    op_ws_close([s.id, closeCode, (reason as string | undefined) ?? null]);
  }
  send(data: unknown): void {
    const s = state(sockets, this);
    required(arguments.length, 1);
    let text = true,
      bytes: Uint8Array<ArrayBuffer>;
    const blob = blobBody(data);
    if (blob !== null) {
      bytes = blob.bytes;
      text = false;
    } else if (isBufferSource(data)) {
      const b = isView(data)
        ? tag(data) === undefined ? dataViewBuffer(data) : buffer(data)
        : data;
      // The getter's brand check also rejects SharedArrayBuffers, which BufferSource excludes.
      if (resizable(b)) {
        throw new TypeError("WebSocket data must have a fixed buffer");
      }
      bytes = snapshot(data);
      text = false;
    } else {
      // The IDL union's last member: everything else becomes a USVString.
      bytes = encode(encoder, usv(data));
    }
    if (s.state === CONNECTING) {
      throw createException({
        message: "WebSocket is not connected",
        name: "InvalidStateError",
      });
    }
    const size = length(bytes);
    s.buffered += size;
    if (s.state !== OPEN) return;
    markPromiseHandled(
      then(op_ws_send([s.id, text, bytes]), (sent) => {
        if (sent && s.state !== CLOSED) s.buffered -= size;
      }),
    );
  }
}

for (
  const [name, value] of [
    ["CONNECTING", 0],
    ["OPEN", 1],
    ["CLOSING", 2],
    ["CLOSED", 3],
  ] as const
) {
  define(WebSocket, name, { value, enumerable: true });
  define(WebSocket.prototype, name, { value, enumerable: true });
}
for (const type of eventTypes) {
  define(WebSocket.prototype, `on${type}`, handler(type));
}

export { adoptWebSocket, CloseEvent, WebSocket };