import * as z from "zod/mini";

const MAX_CHAT_REALTIME_TOKEN_LENGTH = 16_384;
const MAX_CHAT_REALTIME_AUDIENCE_REVISION_LENGTH = 255;

const registrationSequenceSchema = z.int().check(z.positive());
const tokenSchema = z.string().check(
	z.minLength(1),
	z.maxLength(MAX_CHAT_REALTIME_TOKEN_LENGTH),
	z.refine((value) => value === value.trim()),
);
const audienceRevisionSchema = z.string().check(
	z.minLength(1),
	z.maxLength(MAX_CHAT_REALTIME_AUDIENCE_REVISION_LENGTH),
	z.refine((value) => value === value.trim()),
);
const registrationExpiresAtSchema = z.iso.datetime({ precision: 3 });
const registrationExpiresInMsSchema = z.int().check(z.positive());
const jsonObjectSchema = z.record(z.string(), z.unknown());

const chatRealtimeConnectMessageSchema = z.object({
	registrationSequence: registrationSequenceSchema,
	token: tokenSchema,
	type: z.literal("chat.connect"),
	version: z.literal(1),
});

const chatRealtimeConnectedEventSchema = z.object({
	audienceRevision: audienceRevisionSchema,
	registrationExpiresAt: registrationExpiresAtSchema,
	registrationExpiresInMs: registrationExpiresInMsSchema,
	registrationSequence: registrationSequenceSchema,
	type: z.literal("chat.connected"),
	version: z.literal(1),
});

const chatRealtimeRejectedEventSchema = z.discriminatedUnion("retryable", [
	z.object({
		code: z.enum(["identity_unavailable", "service_unavailable"]),
		registrationSequence: registrationSequenceSchema,
		retryable: z.literal(true),
		type: z.literal("chat.rejected"),
		version: z.literal(1),
	}),
	z.object({
		code: z.enum(["invalid_auth", "invalid_request", "superseded"]),
		registrationSequence: registrationSequenceSchema,
		retryable: z.literal(false),
		type: z.literal("chat.rejected"),
		version: z.literal(1),
	}),
]);

const conversationsSyncRealtimeEventSchema = z.object({
	audienceRevision: audienceRevisionSchema,
	type: z.literal("conversations.sync"),
	version: z.literal(1),
});

const chatRealtimeServerEventSchema = z.union([
	chatRealtimeConnectedEventSchema,
	chatRealtimeRejectedEventSchema,
	conversationsSyncRealtimeEventSchema,
]);

const chatRealtimeServerEventTypeSchema = z.enum([
	"chat.connected",
	"chat.rejected",
	"conversations.sync",
]);

export type ChatRealtimeConnectMessage = z.output<
	typeof chatRealtimeConnectMessageSchema
>;
export type ChatRealtimeConnectedEvent = z.output<
	typeof chatRealtimeConnectedEventSchema
>;
export type ChatRealtimeRejectedEvent = z.output<
	typeof chatRealtimeRejectedEventSchema
>;
export type ConversationsSyncRealtimeEvent = z.output<
	typeof conversationsSyncRealtimeEventSchema
>;

export type ChatRealtimeServerEvent =
	| ChatRealtimeConnectedEvent
	| ChatRealtimeRejectedEvent
	| ConversationsSyncRealtimeEvent;

export type ChatRealtimeConnectParseErrorReason =
	| "ignored_message_type"
	| "invalid_json"
	| "invalid_message"
	| "invalid_registration_sequence"
	| "invalid_token"
	| "missing_body"
	| "unsupported_version";

export type ChatRealtimeServerEventParseErrorReason =
	| "ignored_message_type"
	| "invalid_audience_revision"
	| "invalid_json"
	| "invalid_message"
	| "invalid_registration_expiration"
	| "invalid_registration_sequence"
	| "invalid_rejection"
	| "missing_body"
	| "unsupported_data"
	| "unsupported_version";

export type SafeParseChatRealtimeConnectMessageResult =
	| { ok: true; value: ChatRealtimeConnectMessage }
	| {
			ok: false;
			reason: ChatRealtimeConnectParseErrorReason;
			registrationSequence: number | null;
	  };

export type SafeParseChatRealtimeServerEventResult =
	| { ok: true; value: ChatRealtimeServerEvent }
	| { ok: false; reason: ChatRealtimeServerEventParseErrorReason };

export const serializeChatRealtimeConnectMessage = ({
	registrationSequence,
	token,
}: Pick<
	ChatRealtimeConnectMessage,
	"registrationSequence" | "token"
>): string =>
	JSON.stringify({
		registrationSequence: requireRegistrationSequence(registrationSequence),
		token: requireToken(token),
		type: "chat.connect",
		version: 1,
	} satisfies ChatRealtimeConnectMessage);

export const serializeChatRealtimeServerEvent = (
	event: ChatRealtimeServerEvent,
): string => {
	const parsed = chatRealtimeServerEventSchema.safeParse(event);
	if (!parsed.success) {
		throw new TypeError("event must match the chat realtime server protocol.");
	}
	return JSON.stringify(parsed.data);
};

export const safeParseChatRealtimeConnectMessage = (
	body: string | null | undefined,
): SafeParseChatRealtimeConnectMessageResult => {
	const parsed = safeParseJsonObject(body);
	if (!parsed.ok) {
		return { ...parsed, registrationSequence: null };
	}
	if (parsed.value.type !== "chat.connect") {
		return {
			ok: false,
			reason: "ignored_message_type",
			registrationSequence: null,
		};
	}

	const registrationSequence = registrationSequenceSchema.safeParse(
		parsed.value.registrationSequence,
	);
	if (!registrationSequence.success) {
		return {
			ok: false,
			reason: "invalid_registration_sequence",
			registrationSequence: null,
		};
	}
	if (parsed.value.version !== 1) {
		return {
			ok: false,
			reason: "unsupported_version",
			registrationSequence: registrationSequence.data,
		};
	}

	const message = chatRealtimeConnectMessageSchema.safeParse(parsed.value);
	return message.success
		? { ok: true, value: message.data }
		: {
				ok: false,
				reason: "invalid_token",
				registrationSequence: registrationSequence.data,
			};
};

export const safeParseChatRealtimeServerEvent = (
	body: unknown,
): SafeParseChatRealtimeServerEventResult => {
	if (body !== null && body !== undefined && typeof body !== "string") {
		return { ok: false, reason: "unsupported_data" };
	}
	const parsed = safeParseJsonObject(body);
	if (!parsed.ok) {
		return parsed;
	}
	if (!chatRealtimeServerEventTypeSchema.safeParse(parsed.value.type).success) {
		return { ok: false, reason: "ignored_message_type" };
	}
	if (parsed.value.version !== 1) {
		return { ok: false, reason: "unsupported_version" };
	}

	switch (parsed.value.type) {
		case "chat.connected":
			return safeParseConnectedEvent(parsed.value);
		case "chat.rejected": {
			const event = chatRealtimeRejectedEventSchema.safeParse(parsed.value);
			return event.success
				? { ok: true, value: event.data }
				: { ok: false, reason: "invalid_rejection" };
		}
		case "conversations.sync":
			return safeParseConversationsSyncEvent(parsed.value);
		default:
			return { ok: false, reason: "ignored_message_type" };
	}
};

const safeParseConnectedEvent = (
	value: Record<string, unknown>,
): SafeParseChatRealtimeServerEventResult => {
	if (!audienceRevisionSchema.safeParse(value.audienceRevision).success) {
		return { ok: false, reason: "invalid_audience_revision" };
	}
	if (
		!registrationSequenceSchema.safeParse(value.registrationSequence).success
	) {
		return { ok: false, reason: "invalid_registration_sequence" };
	}
	if (
		!registrationExpiresAtSchema.safeParse(value.registrationExpiresAt).success
	) {
		return { ok: false, reason: "invalid_registration_expiration" };
	}
	if (
		!registrationExpiresInMsSchema.safeParse(value.registrationExpiresInMs)
			.success
	) {
		return { ok: false, reason: "invalid_registration_expiration" };
	}
	const event = chatRealtimeConnectedEventSchema.safeParse(value);
	return event.success
		? { ok: true, value: event.data }
		: { ok: false, reason: "invalid_message" };
};

const safeParseConversationsSyncEvent = (
	value: Record<string, unknown>,
): SafeParseChatRealtimeServerEventResult => {
	if (!audienceRevisionSchema.safeParse(value.audienceRevision).success) {
		return { ok: false, reason: "invalid_audience_revision" };
	}
	const event = conversationsSyncRealtimeEventSchema.safeParse(value);
	return event.success
		? { ok: true, value: event.data }
		: { ok: false, reason: "invalid_message" };
};

type ChatRealtimeObjectParseErrorReason =
	| "invalid_json"
	| "invalid_message"
	| "missing_body";

const safeParseJsonObject = (
	body: string | null | undefined,
):
	| { ok: true; value: Record<string, unknown> }
	| { ok: false; reason: ChatRealtimeObjectParseErrorReason } => {
	if (!body) {
		return { ok: false, reason: "missing_body" };
	}
	let parsed: unknown;
	try {
		parsed = JSON.parse(body) as unknown;
	} catch {
		return { ok: false, reason: "invalid_json" };
	}
	const object = jsonObjectSchema.safeParse(parsed);
	return object.success
		? { ok: true, value: object.data }
		: { ok: false, reason: "invalid_message" };
};

const requireRegistrationSequence = (value: unknown): number => {
	const parsed = registrationSequenceSchema.safeParse(value);
	if (!parsed.success) {
		throw new TypeError(
			"registrationSequence must be a positive safe integer.",
		);
	}
	return parsed.data;
};

const requireToken = (value: unknown): string => {
	const parsed = tokenSchema.safeParse(value);
	if (!parsed.success) {
		throw new TypeError(
			`token must be a trimmed non-empty string of at most ${MAX_CHAT_REALTIME_TOKEN_LENGTH} characters.`,
		);
	}
	return parsed.data;
};