import type {
  BinaryCodec,
  ClientStateReport,
  InputPacket,
  PingPacket,
  PlayerId,
  PongPacket,
  StateSnapshot,
} from "./types.js";

const enum ServerOpcode {
  Welcome = 1,
  Snapshot = 2,
  Acknowledge = 3,
  Pong = 4,
  Validation = 5,
  RejectInput = 6,
  Event = 7,
  ReplayStart = 8,
  ReplayFrame = 9,
  ReplayEnd = 10,
}

const enum ClientOpcode {
  Input = 16,
  Ping = 17,
  StateReport = 18,
}

export type ClientWireMessage<Input, State> =
  | { kind: "input"; packet: InputPacket<Input> }
  | { kind: "ping"; ping: PingPacket }
  | { kind: "state-report"; report: ClientStateReport<State> };

export type ServerWireMessage<State, Event = never> =
  | { kind: "welcome"; playerId: PlayerId; snapshot: StateSnapshot<State> }
  | { kind: "snapshot"; snapshot: StateSnapshot<State> }
  | { kind: "acknowledge"; sequence: number }
  | { kind: "pong"; pong: PongPacket }
  | { kind: "validation"; tick: number; valid: boolean }
  | { kind: "reject-input"; sequence: number }
  | { kind: "event"; tick: number; event: Event }
  | {
      kind: "replay-start";
      ticketId: number;
      perspectiveId: PlayerId;
      fromTick: number;
      toTick: number;
      frameCount: number;
      playbackRate: number;
    }
  | {
      kind: "replay-frame";
      ticketId: number;
      tick: number;
      state: State;
      events: Event[];
    }
  | { kind: "replay-end"; ticketId: number };

export interface BinaryProtocol<Input, State, Event = never> {
  encodeClient(message: ClientWireMessage<Input, State>): ArrayBuffer;
  decodeClient(payload: ArrayBuffer): ClientWireMessage<Input, State>;
  encodeServer(message: ServerWireMessage<State, Event>): ArrayBuffer;
  decodeServer(payload: ArrayBuffer): ServerWireMessage<State, Event>;
}

export function createBinaryProtocol<Input, State, Event = never>(codecs: {
  input: BinaryCodec<Input>;
  state: BinaryCodec<State>;
  event?: BinaryCodec<Event>;
}): BinaryProtocol<Input, State, Event> {
  return {
    encodeClient(message) {
      switch (message.kind) {
        case "input": {
          const input = codecs.input.encode(message.packet.input);
          const payload = createFrame(ClientOpcode.Input, 13, input);
          const view = new DataView(payload);
          view.setUint32(1, message.packet.sequence, true);
          view.setUint32(5, message.packet.targetTick, true);
          view.setUint32(
            9,
            message.packet.observedTick ?? message.packet.targetTick,
            true,
          );
          return payload;
        }
        case "ping": {
          const payload = createFrame(ClientOpcode.Ping, 13);
          const view = new DataView(payload);
          view.setUint32(1, message.ping.id, true);
          view.setFloat64(5, message.ping.clientSentAt, true);
          return payload;
        }
        case "state-report": {
          const state = codecs.state.encode(message.report.state);
          const payload = createFrame(ClientOpcode.StateReport, 5, state);
          new DataView(payload).setUint32(1, message.report.tick, true);
          return payload;
        }
      }
    },
    decodeClient(payload) {
      const view = new DataView(payload);
      switch (view.getUint8(0)) {
        case ClientOpcode.Input:
          assertLength(payload, 13);
          return {
            kind: "input",
            packet: {
              sequence: view.getUint32(1, true),
              targetTick: view.getUint32(5, true),
              observedTick: view.getUint32(9, true),
              input: codecs.input.decode(new Uint8Array(payload, 13)),
            },
          };
        case ClientOpcode.Ping:
          assertExactLength(payload, 13);
          return {
            kind: "ping",
            ping: {
              id: view.getUint32(1, true),
              clientSentAt: view.getFloat64(5, true),
            },
          };
        case ClientOpcode.StateReport:
          assertLength(payload, 5);
          return {
            kind: "state-report",
            report: {
              tick: view.getUint32(1, true),
              state: codecs.state.decode(new Uint8Array(payload, 5)),
            },
          };
        default:
          throw new RangeError("Unknown client message opcode");
      }
    },
    encodeServer(message) {
      switch (message.kind) {
        case "welcome": {
          const state = codecs.state.encode(message.snapshot.state);
          const payload = createFrame(ServerOpcode.Welcome, 17, state);
          const view = new DataView(payload);
          view.setUint32(1, message.playerId, true);
          view.setUint32(5, message.snapshot.tick, true);
          view.setFloat64(9, message.snapshot.serverTime, true);
          return payload;
        }
        case "snapshot": {
          const state = codecs.state.encode(message.snapshot.state);
          const payload = createFrame(ServerOpcode.Snapshot, 13, state);
          const view = new DataView(payload);
          view.setUint32(1, message.snapshot.tick, true);
          view.setFloat64(5, message.snapshot.serverTime, true);
          return payload;
        }
        case "acknowledge": {
          const payload = createFrame(ServerOpcode.Acknowledge, 5);
          new DataView(payload).setUint32(1, message.sequence, true);
          return payload;
        }
        case "pong": {
          const payload = createFrame(ServerOpcode.Pong, 29);
          const view = new DataView(payload);
          view.setUint32(1, message.pong.id, true);
          view.setFloat64(5, message.pong.clientSentAt, true);
          view.setFloat64(13, message.pong.serverReceivedAt, true);
          view.setFloat64(21, message.pong.serverSentAt, true);
          return payload;
        }
        case "validation": {
          const payload = createFrame(ServerOpcode.Validation, 6);
          const view = new DataView(payload);
          view.setUint32(1, message.tick, true);
          view.setUint8(5, message.valid ? 1 : 0);
          return payload;
        }
        case "reject-input": {
          const payload = createFrame(ServerOpcode.RejectInput, 5);
          new DataView(payload).setUint32(1, message.sequence, true);
          return payload;
        }
        case "event": {
          if (!codecs.event) {
            throw new Error("This protocol has no event codec");
          }
          const event = codecs.event.encode(message.event);
          const payload = createFrame(ServerOpcode.Event, 5, event);
          new DataView(payload).setUint32(1, message.tick, true);
          return payload;
        }
        case "replay-start": {
          if (message.frameCount > 65_535) {
            throw new RangeError("Replay frame count exceeds protocol limit");
          }
          const payload = createFrame(ServerOpcode.ReplayStart, 23);
          const view = new DataView(payload);
          view.setUint32(1, message.ticketId, true);
          view.setUint32(5, message.perspectiveId, true);
          view.setUint32(9, message.fromTick, true);
          view.setUint32(13, message.toTick, true);
          view.setUint16(17, message.frameCount, true);
          view.setFloat32(19, message.playbackRate, true);
          return payload;
        }
        case "replay-frame": {
          if (message.events.length > 65_535) {
            throw new RangeError("Replay event count exceeds protocol limit");
          }
          if (message.events.length > 0 && !codecs.event) {
            throw new Error("This protocol has no event codec");
          }
          const state = codecs.state.encode(message.state);
          const events = message.events.map((event) => codecs.event!.encode(event));
          const bodyLength =
            state.byteLength +
            events.reduce((total, event) => total + 4 + event.byteLength, 0);
          const payload = createFrame(ServerOpcode.ReplayFrame, 15 + bodyLength);
          const view = new DataView(payload);
          view.setUint32(1, message.ticketId, true);
          view.setUint32(5, message.tick, true);
          view.setUint16(9, events.length, true);
          view.setUint32(11, state.byteLength, true);
          const bytes = new Uint8Array(payload);
          let offset = 15;
          bytes.set(state, offset);
          offset += state.byteLength;
          for (const event of events) {
            view.setUint32(offset, event.byteLength, true);
            offset += 4;
            bytes.set(event, offset);
            offset += event.byteLength;
          }
          return payload;
        }
        case "replay-end": {
          const payload = createFrame(ServerOpcode.ReplayEnd, 5);
          new DataView(payload).setUint32(1, message.ticketId, true);
          return payload;
        }
      }
    },
    decodeServer(payload) {
      const view = new DataView(payload);
      switch (view.getUint8(0)) {
        case ServerOpcode.Welcome:
          assertLength(payload, 17);
          return {
            kind: "welcome",
            playerId: view.getUint32(1, true),
            snapshot: {
              tick: view.getUint32(5, true),
              serverTime: view.getFloat64(9, true),
              state: codecs.state.decode(new Uint8Array(payload, 17)),
            },
          };
        case ServerOpcode.Snapshot:
          assertLength(payload, 13);
          return {
            kind: "snapshot",
            snapshot: {
              tick: view.getUint32(1, true),
              serverTime: view.getFloat64(5, true),
              state: codecs.state.decode(new Uint8Array(payload, 13)),
            },
          };
        case ServerOpcode.Acknowledge:
          assertExactLength(payload, 5);
          return {
            kind: "acknowledge",
            sequence: view.getUint32(1, true),
          };
        case ServerOpcode.Pong:
          assertExactLength(payload, 29);
          return {
            kind: "pong",
            pong: {
              id: view.getUint32(1, true),
              clientSentAt: view.getFloat64(5, true),
              serverReceivedAt: view.getFloat64(13, true),
              serverSentAt: view.getFloat64(21, true),
            },
          };
        case ServerOpcode.Validation:
          assertExactLength(payload, 6);
          return {
            kind: "validation",
            tick: view.getUint32(1, true),
            valid: view.getUint8(5) === 1,
          };
        case ServerOpcode.RejectInput:
          assertExactLength(payload, 5);
          return {
            kind: "reject-input",
            sequence: view.getUint32(1, true),
          };
        case ServerOpcode.Event:
          assertLength(payload, 5);
          if (!codecs.event) {
            throw new Error("This protocol has no event codec");
          }
          return {
            kind: "event",
            tick: view.getUint32(1, true),
            event: codecs.event.decode(new Uint8Array(payload, 5)),
          };
        case ServerOpcode.ReplayStart:
          assertExactLength(payload, 23);
          return {
            kind: "replay-start",
            ticketId: view.getUint32(1, true),
            perspectiveId: view.getUint32(5, true),
            fromTick: view.getUint32(9, true),
            toTick: view.getUint32(13, true),
            frameCount: view.getUint16(17, true),
            playbackRate: view.getFloat32(19, true),
          };
        case ServerOpcode.ReplayFrame: {
          assertLength(payload, 15);
          const ticketId = view.getUint32(1, true);
          const tick = view.getUint32(5, true);
          const eventCount = view.getUint16(9, true);
          const stateLength = view.getUint32(11, true);
          assertAvailable(payload, 15, stateLength);
          const state = codecs.state.decode(
            new Uint8Array(payload, 15, stateLength),
          );
          let offset = 15 + stateLength;
          const events: Event[] = [];
          for (let index = 0; index < eventCount; index += 1) {
            if (!codecs.event) throw new Error("This protocol has no event codec");
            assertAvailable(payload, offset, 4);
            const eventLength = view.getUint32(offset, true);
            offset += 4;
            assertAvailable(payload, offset, eventLength);
            events.push(
              codecs.event.decode(new Uint8Array(payload, offset, eventLength)),
            );
            offset += eventLength;
          }
          if (offset !== payload.byteLength) {
            throw new RangeError("Replay frame has trailing bytes");
          }
          return { kind: "replay-frame", ticketId, tick, state, events };
        }
        case ServerOpcode.ReplayEnd:
          assertExactLength(payload, 5);
          return {
            kind: "replay-end",
            ticketId: view.getUint32(1, true),
          };
        default:
          throw new RangeError("Unknown server message opcode");
      }
    },
  };
}

function createFrame(
  opcode: number,
  headerLength: number,
  body?: Uint8Array,
): ArrayBuffer {
  const payload = new ArrayBuffer(headerLength + (body?.byteLength ?? 0));
  const bytes = new Uint8Array(payload);
  bytes[0] = opcode;
  if (body) {
    bytes.set(body, headerLength);
  }
  return payload;
}

function assertLength(payload: ArrayBuffer, minimum: number): void {
  if (payload.byteLength < minimum) {
    throw new RangeError(`Message must be at least ${minimum} bytes`);
  }
}

function assertExactLength(payload: ArrayBuffer, expected: number): void {
  if (payload.byteLength !== expected) {
    throw new RangeError(`Message must be exactly ${expected} bytes`);
  }
}

function assertAvailable(
  payload: ArrayBuffer,
  offset: number,
  length: number,
): void {
  if (
    !Number.isInteger(offset) ||
    !Number.isInteger(length) ||
    offset < 0 ||
    length < 0 ||
    offset > payload.byteLength - length
  ) {
    throw new RangeError("Message body is truncated");
  }
}