import { jwtVerify } from "@panva/jose";
import { combineSignals, timeoutSignal } from "@repo/common/aborts.ts";
import { jwtDecodedSchema } from "@repo/common/crypto-sig.ts";
import { AuthError, normalizeError } from "@repo/common/errors.ts";
import { takeSocket } from "@repo/common/socket.ts";
import { StrictMap } from "@repo/common/strict-map.ts";
import { preauthAuthnMessageSchema } from "@repo/schema/proto.ts";
import { parseIdentId, parseRealmId } from "@repo/schema/state.ts";

import { Realm, realmMap, Session } from "./state.ts";

export async function preauthHandler(
  ws: WebSocket,
  signal: AbortSignal,
): Promise<Session & { realm: Realm }> {
  const timeout = timeoutSignal(3000);
  const combinedSignal = combineSignals(signal, timeout.signal);

  // wait for preauth.authn message;it will be signed by the identity,
  // which we already have if the realm exists, and we add to a new realm if being created
  try {
    const data = await takeSocket(ws, combinedSignal);
    const jwt = jwtDecodedSchema.parse(data);

    const msg = await preauthAuthnMessageSchema.parseAsync(jwt);
    const realmid = parseRealmId(jwt.aud);
    const identid = parseIdentId(jwt.iss);

    // make sure we have a realm
    // if a new one is being created, the current identity to it
    const realm = realmMap.ensure(realmid, () => ({
      id: realmid,
      sockets: new StrictMap(),
      identities: new StrictMap([[identid, msg.pubkey]]),
    }));

    try {
      // we've looked up the realm
      // pubkey should match what we've got stored, and the signature should be valid
      const pubkey = realm.identities.require(identid);
      await jwtVerify(data, pubkey);

      // the signature is good
      return { realm, realmid, identid, pubkey };
    } catch (e) {
      throw new AuthError("jwt verification failed", normalizeError(e));
    }
  } finally {
    timeout.cleanup();
  }
}