import "./go.js";

import type {
  DescMessage,
  MessageShape,
  MessageValidType,
} from "@bufbuild/protobuf";
import { fromBinary as bufFromBinary, toBinary } from "@bufbuild/protobuf";
import { AnySchema, anyIs } from "@bufbuild/protobuf/wkt";

import type {
  CheckResourcesResponse,
  Config,
  PlanResourcesResponse,
} from "@cerbos/api/cerbos/cloud/epdp/v2/epdp_pb";
import {
  CheckResourcesResponseSchema,
  ConfigSchema,
  ErrorSchema,
  PlanResourcesResponseSchema,
} from "@cerbos/api/cerbos/cloud/epdp/v2/epdp_pb";
import type { AuxDataValid } from "@cerbos/api/cerbos/engine/v1/engine_pb";
import { AuxDataSchema } from "@cerbos/api/cerbos/engine/v1/engine_pb";
import type {
  CheckResourcesRequestValid,
  PlanResourcesRequestValid,
} from "@cerbos/api/cerbos/request/v1/request_pb";
import {
  CheckResourcesRequestSchema,
  PlanResourcesRequestSchema,
} from "@cerbos/api/cerbos/request/v1/request_pb";
import { Code } from "@cerbos/api/google/rpc/code_pb";

import { ServerError } from "./error.js";

export type Instantiate = (
  imports: WebAssembly.Imports,
) => Promise<WebAssembly.Instance | WebAssembly.WebAssemblyInstantiatedSource>;

interface Instance extends WebAssembly.Instance {
  exports: Exports;
}

interface Exports extends WebAssembly.Exports {
  mem: WebAssembly.Memory;
  alloc(length: number): bigint;
  free(pointer: bigint): void;
  configure(pointer: bigint): bigint;
  loadRuleTable(pointer: bigint): bigint;
  checkResources(requestPointer: bigint, auxDataPointer: bigint): bigint;
  planResources(requestPointer: bigint, auxDataPointer: bigint): bigint;
}

export class Server {
  private exited?: ServerError;

  public static async from(
    instantiate: Instantiate,
    config: Config | undefined,
  ): Promise<Server> {
    const go = new Go();
    const instantiated = await instantiate(go.importObject);

    const instance = (
      instantiated instanceof WebAssembly.Instance
        ? instantiated
        : instantiated.instance
    ) as Instance;

    return new Server(go, instance, config);
  }

  private constructor(
    go: Go,
    private readonly instance: Instance,
    config: Config | undefined,
  ) {
    go.run(instance).then(
      () => {
        this.exited = new ServerError(
          Code.INTERNAL,
          "embedded server has exited",
        );
      },
      (error: unknown) => {
        this.exited = new ServerError(
          Code.INTERNAL,
          error instanceof Error ? error.message : "embedded server has failed",
          { cause: error },
        );
      },
    );

    if (config) {
      this.configure(config);
    }
  }

  public loadRuleTable(ruleTable: Uint8Array): void {
    const pointer = this.write(ruleTable);
    try {
      const result = this.exports.loadRuleTable(pointer);
      if (result !== 0n) {
        this.decode(ErrorSchema, result);
      }
    } finally {
      this.free(pointer);
    }
  }

  public checkResources(
    request: CheckResourcesRequestValid,
    auxData: AuxDataValid | undefined,
  ): CheckResourcesResponse {
    const requestPointer = this.encode(CheckResourcesRequestSchema, request);
    try {
      const auxDataPointer = this.encode(AuxDataSchema, auxData);
      try {
        return this.decode(
          CheckResourcesResponseSchema,
          this.exports.checkResources(requestPointer, auxDataPointer),
        );
      } finally {
        this.free(auxDataPointer);
      }
    } finally {
      this.free(requestPointer);
    }
  }

  public planResources(
    request: PlanResourcesRequestValid,
    auxData: AuxDataValid | undefined,
  ): PlanResourcesResponse {
    const requestPointer = this.encode(PlanResourcesRequestSchema, request);
    try {
      const auxDataPointer = this.encode(AuxDataSchema, auxData);
      try {
        return this.decode(
          PlanResourcesResponseSchema,
          this.exports.planResources(requestPointer, auxDataPointer),
        );
      } finally {
        this.free(auxDataPointer);
      }
    } finally {
      this.free(requestPointer);
    }
  }

  private configure(config: Config): void {
    const pointer = this.encode(ConfigSchema, config);
    try {
      const result = this.exports.configure(pointer);
      if (result !== 0n) {
        this.decode(ErrorSchema, result);
      }
    } finally {
      this.free(pointer);
    }
  }

  private decode<T extends DescMessage>(
    schema: T,
    pointer: bigint,
  ): T extends typeof ErrorSchema ? never : MessageShape<T>;

  private decode<T extends DescMessage>(
    schema: T,
    pointer: bigint,
  ): MessageShape<T> {
    try {
      const message = fromBinary(AnySchema, this.slice(pointer));

      if (anyIs(message, ErrorSchema)) {
        throw ServerError.from(fromBinary(ErrorSchema, message.value));
      }

      if (!anyIs(message, schema)) {
        throw new ServerError(
          Code.INTERNAL,
          `unexpected response type "${message.typeUrl}"`,
        );
      }

      return fromBinary(schema, message.value);
    } finally {
      this.free(pointer);
    }
  }

  private encode<T extends DescMessage>(
    schema: T,
    message: MessageValidType<T> | undefined,
  ): bigint {
    if (!message) {
      return 0n;
    }

    return this.write(toBinary(schema, message as MessageShape<T>));
  }

  private write(data: Uint8Array): bigint {
    const pointer = this.exports.alloc(data.length);
    this.slice(pointer).set(data);
    return pointer;
  }

  private slice(pointer: bigint): Uint8Array {
    return new Uint8Array(
      this.exports.mem.buffer,
      Number(pointer >> 32n),
      Number(pointer & 0xffffffffn),
    );
  }

  private get exports(): Exports {
    if (this.exited) {
      throw this.exited;
    }

    return this.instance.exports;
  }

  private free(pointer: bigint): void {
    if (!this.exited && pointer !== 0n) {
      this.exports.free(pointer);
    }
  }
}

function fromBinary<T extends DescMessage>(
  schema: T,
  bytes: Uint8Array,
): MessageShape<T> {
  return bufFromBinary(schema, bytes, { readUnknownFields: false });
}
