import { jwkSchema } from "@repo/common/crypto-sig.ts";
import { z } from "zod/v4";

import { identIdSchema, realmIdSchema } from "./state.ts";

export const preauthAuthnMessageSchema = z.object({
  msg: z.literal("preauth.authn"),
  pubkey: jwkSchema,
});

export type PreauthAuthnMessage = z.infer<typeof preauthAuthnMessageSchema>;

export const preauthMessageSchema = z.discriminatedUnion("msg", [
  preauthAuthnMessageSchema,
]);

export type PreauthMessage = z.infer<typeof preauthMessageSchema>;

export const realmStatusMessageSchema = z.object({
  msg: z.literal("realm.status"),
});

export const realmStatusResponseSchema = z.object({
  msg: z.literal("realm.status"),
  realm: realmIdSchema,
  identities: z.array(identIdSchema),
});

export type RealmStatusMessage = z.infer<typeof realmStatusMessageSchema>;
export type RealmStatusResponse = z.infer<typeof realmStatusResponseSchema>;

export const realmBroadcastMessageSchema = z.object({
  msg: z.literal("realm.broadcast"),
  payload: z.any().nonoptional(),
  recipients: z.array(identIdSchema),
});

export type RealmBroadcastMessage = z.infer<typeof realmBroadcastMessageSchema>;

export const realmMessageSchema = z.discriminatedUnion("msg", [
  realmStatusMessageSchema,
  realmBroadcastMessageSchema,
]);

export type RealmMessage = z.infer<typeof realmMessageSchema>;