import WebSocket from 'isomorphic-ws'

import {combineSignals, timeoutSignal} from '#common/async/aborts'
import {JWK, jwkExport, jwkImport} from '#common/crypto/jwks'
import {jwtPayload, jwtSchema, JWTToken, verifyJwtToken} from '#common/crypto/jwts'
import {normalizeError, ProtocolError} from '#common/errors'
import {takeSocket} from '#common/socket'
import {IdentBrand, IdentID, preauthReqSchema, PreauthResponse, RealmBrand, RealmID} from '#realm/protocol/index'

import * as realms from './state'

/**
 * immediately after the socket connects, we up to 3 seconds for a valid authentication message.
 * - if the realm does not exist (by realm id), we create a new one, and add the identity (success).
 * - if the realm /does/ exist, we verify the message is a signed JWT our already registered pubkey.
 */
export async function preauthHandler(ws: WebSocket, signal?: AbortSignal): Promise<realms.AuthenticatedIdentity> {
  const timeout = timeoutSignal(3000)
  const combinedSignal = combineSignals(signal, timeout.signal)

  try {
    const data = await takeSocket(ws, combinedSignal)
    const jwt = await jwtPayload(preauthReqSchema).parseAsync(data)

    // if any of the parsing fails, it'll throw a zod error
    const identid = IdentBrand.parse(jwt.claims.iss)
    const realmid = RealmBrand.parse(jwt.claims.aud)

    switch (jwt.payload.msg) {
      case 'preauth.register': {
        // if we're registering, make sure the realm exists
        const registrantkey = await jwkImport.parseAsync(jwt.payload.dat.pubkey)
        await realms.ensureRegisteredRealm(realmid, identid, registrantkey)

        break
      }

      case 'preauth.exchange': {
        // validate and then insert (validation throws)
        const token = jwtSchema.parse(jwt.payload.dat.inviteJwt)
        await preauthValidateInvitation(realmid, token)

        const inviteeid = IdentBrand.parse(jwt.claims.iss)
        const inviteekey = await jwkImport.parseAsync(jwt.payload.dat.pubkey)
        await realms.admitToRealm(realmid, inviteeid, inviteekey)

        break
      }
    }

    // everything is in place, fall through to authentication
    const auth = await authenticatePreauth(realmid, identid, jwt.token)
    const msg = await preauthResponse(auth, jwt.payload.seq)
    ws.send(JSON.stringify(msg))

    return auth
  } finally {
    timeout.cancel()
  }
}

async function preauthValidateInvitation(realmid: RealmID, invitation: JWTToken): Promise<void> {
  try {
    const realmMap = await realms.ensureRealmMap()
    const realm = realmMap.require(realmid)

    if (!invitation.claims.jti) throw new Error('invitation requires nonce!')
    if (!(await realms.validateNonce(realmid, invitation.claims.jti))) throw new Error('invitation already used!')

    const inviterid = IdentBrand.parse(invitation.claims.iss)
    const inviterkey = realm.storage.identities.require(inviterid)
    await verifyJwtToken(invitation.token, inviterkey, {subject: 'invitation'})
  } catch (exc) {
    const err = normalizeError(exc)
    throw new ProtocolError('invitation verification failed', 401, {cause: err})
  }
}

async function authenticatePreauth(
  realmid: RealmID,
  identid: IdentID,
  token: string,
): Promise<realms.AuthenticatedIdentity> {
  try {
    const realmMap = await realms.ensureRealmMap()
    const realm = realmMap.require(realmid)
    const pubkey = realm.storage.identities.require(identid)

    // at this point we no langer care about the payload
    // but this throws as a side-effect if the token is invalid
    await verifyJwtToken(token, pubkey)
    return {realmid, realm, identid, pubkey}
  } catch (exc) {
    const err = normalizeError(exc)
    throw new ProtocolError('jwt verification failed', 401, {cause: err})
  }
}

async function preauthResponse(auth: realms.AuthenticatedIdentity, seq?: number): Promise<PreauthResponse> {
  const peers: Array<IdentID | 'server'> = Array.from(auth.realm.sockets.keys())
  peers.unshift('server')

  const identities: Record<IdentID, JWK> = {}
  for (const identid of auth.realm.storage.identities.keys()) {
    const pubkey = auth.realm.storage.identities.require(identid)
    identities[identid] = await jwkExport.parseAsync(pubkey)
  }

  return {
    typ: 'res',
    msg: 'preauth.authn',
    dat: {peers, identities},
    seq,
  }
}