cloudflare_temp_email/worker/src/mail_flags.ts

export enum MailFlag {
    UNREAD = 0,
    ANSWERED = 1,
    FLAGGED = 2,
    DELETED = 3,
    DRAFT = 4,
    JUNK = 5,
}

export const CUSTOM_MAIL_FLAG_OFFSET = 10;
export const CUSTOM_MAIL_FLAG_COUNT = 10;

export enum MailState {
    ALL = 'all',
    UNREAD = 'unread',
    READ = 'read',
}

export type MailStateOption = {
    value: string;
    label_key?: string;
    label?: string;
    unread?: boolean;
    default?: boolean;
};

type MailStateMutation = {
    add?: number;
    remove?: number[];
};

export type MailStateDefinition = MailStateOption & {
    filter?: { flag: number; present: boolean };
    mutation?: MailStateMutation;
};

const SYSTEM_MAIL_STATES: MailStateDefinition[] = [
    { value: MailState.ALL, label_key: 'allMail', default: true },
    {
        value: MailState.UNREAD,
        label_key: 'unread',
        unread: true,
        filter: { flag: MailFlag.UNREAD, present: true },
        mutation: { add: MailFlag.UNREAD },
    },
    {
        value: MailState.READ,
        label_key: 'read',
        unread: false,
        filter: { flag: MailFlag.UNREAD, present: false },
        mutation: { remove: [MailFlag.UNREAD] },
    },
];

export const getMailStateOptions = (
    customStates: MailStateDefinition[] = [],
): MailStateOption[] => {
    return [...SYSTEM_MAIL_STATES, ...customStates].map(state => {
        const { filter: _filter, mutation: _mutation, ...option } = state;
        return option;
    });
};

const getMailStateDefinition = (
    value: unknown,
    customStates: MailStateDefinition[] = [],
) => {
    if (typeof value !== 'string') return undefined;
    return [...SYSTEM_MAIL_STATES, ...customStates].find(state => state.value === value);
};

export const getCustomMailFlag = (slot: number): number => {
    if (!Number.isInteger(slot) || slot < 0 || slot >= CUSTOM_MAIL_FLAG_COUNT) {
        throw new Error('Invalid custom mail flag slot');
    }
    return CUSTOM_MAIL_FLAG_OFFSET + slot;
};

export type CustomMailStateConfig = {
    slot: number;
    name: string;
};

const CUSTOM_MAIL_FLAGS = Array.from(
    { length: CUSTOM_MAIL_FLAG_COUNT },
    (_, slot) => getCustomMailFlag(slot),
);

export const createCustomMailStateDefinitions = (
    configs: CustomMailStateConfig[],
): MailStateDefinition[] => {
    return configs.map(config => {
        const flag = getCustomMailFlag(config.slot);
        return {
            value: `custom:${config.slot}`,
            label: config.name,
            filter: { flag, present: true },
            mutation: { add: flag, remove: CUSTOM_MAIL_FLAGS.filter(value => value !== flag) },
        };
    });
};

const withoutLegacyFlags = <T extends Record<string, unknown>>(row: T): T => {
    const result = { ...row };
    delete result.flags;
    return result;
};

const isMailFlagsEnabled = (env: Bindings): boolean => {
    return env.ENABLE_MAIL_FLAGS === true || env.ENABLE_MAIL_FLAGS === 'true';
};

export const serializeMailStates = async <T extends Record<string, unknown>>(
    db: D1Database,
    rows: T[],
    env: Bindings,
): Promise<T[]> => {
    const results = rows.map(withoutLegacyFlags);
    if (!isMailFlagsEnabled(env) || results.length === 0) return results;

    const needsUnread = (row: T) => ![true, false, 0, 1].includes(row.unread as boolean | number);
    const ids = [...new Set(results.filter(needsUnread).map(row => Number(row.id)))]
        .filter(id => Number.isInteger(id) && id > 0);
    if (ids.length === 0) {
        return results.map(row => ({ ...row, unread: Boolean(row.unread) }));
    }

    const placeholders = ids.map(() => '?').join(',');
    const { results: flagRows } = await db.prepare(
        `SELECT mf.mail_id FROM mail_flags mf`
        + ` JOIN raw_mails rm ON rm.id = mf.mail_id`
        + ` JOIN address a ON a.id = mf.address_id AND a.name = rm.address`
        + ` WHERE mf.flag = ? AND mf.mail_id IN (${placeholders})`
    ).bind(MailFlag.UNREAD, ...ids).all<{ mail_id: number }>();
    const unreadIds = new Set(flagRows.map(row => Number(row.mail_id)));

    return results.map(row => ({
        ...row,
        unread: needsUnread(row) ? unreadIds.has(Number(row.id)) : Boolean(row.unread),
    }));
};

export const serializeMailState = async <T extends Record<string, unknown>>(
    db: D1Database,
    row: T,
    env: Bindings,
): Promise<T> => {
    const [result] = await serializeMailStates(db, [row], env);
    return result;
};

const resolveInitialMailFlags = async (
    _env: Bindings,
    _address: string,
    _parsedEmailContext: ParsedEmailContext,
): Promise<number[]> => {
    return [MailFlag.UNREAD];
};

export const updateInitialMailFlags = async (
    db: D1Database,
    enabled: boolean,
    mailId: number,
    env: Bindings,
    address: string,
    parsedEmailContext: ParsedEmailContext,
) => {
    if (!enabled || !Number.isInteger(mailId) || mailId <= 0) return;

    const flags = await resolveInitialMailFlags(env, address, parsedEmailContext);
    if (flags.length === 0) return;

    const selects = flags.map(() => 'SELECT ?, id, ? FROM address WHERE name = ?').join(' UNION ALL ');
    await db.prepare(
        `INSERT OR IGNORE INTO mail_flags (mail_id, address_id, flag) ${selects}`
    ).bind(...flags.flatMap(flag => [mailId, flag, address])).run();
};

export type MailStateQuery = {
    join: string;
    clause?: string;
    orderBy?: string;
    unread?: boolean;
    params: number[];
};

export const getMailStateQuery = (
    value: string | undefined,
    mailAlias: string,
    addressIdColumn: string,
    customStates: MailStateDefinition[] = [],
): MailStateQuery | undefined | null => {
    if (value === undefined) return undefined;

    const definition = getMailStateDefinition(value, customStates);
    if (!definition) return null;
    if (!definition.filter) return undefined;

    const { flag, present } = definition.filter;
    const joinType = present ? 'JOIN' : 'LEFT JOIN';
    return {
        join: ` ${joinType} mail_flags mail_state_flags`
            + ` ON mail_state_flags.mail_id = ${mailAlias}.id`
            + ` AND mail_state_flags.address_id = ${addressIdColumn}`
            + ` AND mail_state_flags.flag = ?`,
        clause: present ? undefined : 'mail_state_flags.mail_id IS NULL',
        orderBy: present ? 'mail_state_flags.mail_id desc' : undefined,
        unread: definition.unread,
        params: [flag],
    };
};

type MailStateUpdate = {
    ids: number[];
    mutation: MailStateMutation;
    unread?: boolean;
};

const parseMailStateUpdate = (
    value: unknown,
    customStates: MailStateDefinition[] = [],
): MailStateUpdate | null => {
    if (!value || typeof value !== 'object') return null;

    const body = value as Record<string, unknown>;
    if (!Array.isArray(body.ids) || body.ids.length === 0 || body.ids.length > 100) return null;
    if (body.ids.some(id => typeof id !== 'number')) return null;

    const ids = [...new Set(body.ids.map(Number))];
    if (ids.some(id => !Number.isInteger(id) || id <= 0)) return null;

    const definition = getMailStateDefinition(body.state, customStates);
    if (!definition?.mutation) return null;
    return { ids, mutation: definition.mutation, unread: definition.unread };
};

type MailScope = {
    clause: string;
    params: (string | number)[];
};

export const applyMailStateUpdate = async (
    db: D1Database,
    env: Bindings,
    scope: MailScope,
    value: unknown,
    customStates: MailStateDefinition[] = [],
) => {
    const update = parseMailStateUpdate(value, customStates);
    if (!update) return null;

    const idPlaceholders = update.ids.map(() => '?').join(',');
    const targetWhere = `rm.id IN (${idPlaceholders}) AND (${scope.clause})`;
    const statements: D1PreparedStatement[] = [];

    if (update.mutation.remove?.length) {
        const flagPlaceholders = update.mutation.remove.map(() => '?').join(',');
        statements.push(db.prepare(
            `DELETE FROM mail_flags WHERE flag IN (${flagPlaceholders})`
            + ` AND mail_id IN (`
            + `SELECT rm.id FROM raw_mails rm JOIN address a ON a.name = rm.address`
            + ` WHERE ${targetWhere})`
        ).bind(...update.mutation.remove, ...update.ids, ...scope.params));
    }

    if (update.mutation.add !== undefined) {
        statements.push(db.prepare(
            `INSERT OR IGNORE INTO mail_flags (mail_id, address_id, flag)`
            + ` SELECT rm.id, a.id, ?`
            + ` FROM raw_mails rm JOIN address a ON a.name = rm.address`
            + ` WHERE ${targetWhere}`
        ).bind(update.mutation.add, ...update.ids, ...scope.params));
    }

    const mutationResults = await db.batch(statements);
    if (mutationResults.some(result => !result.success)) {
        return { success: false, changes: 0, results: [] };
    }

    const unreadSelect = update.unread === undefined ? '' : `, ${update.unread ? 1 : 0} AS unread`;
    const { results } = await db.prepare(
        `SELECT rm.id${unreadSelect} FROM raw_mails rm JOIN address a ON a.name = rm.address`
        + ` WHERE ${targetWhere}`
    ).bind(...update.ids, ...scope.params).all();

    return {
        success: true,
        changes: mutationResults.reduce((total, result) => total + (result.meta.changes ?? 0), 0),
        results: await serializeMailStates(db, results, env),
    };
};

export const prepareRawMailDeleteStatements = (
    db: D1Database,
    env: Bindings,
    whereClause: string,
    params: (string | number)[],
): D1PreparedStatement[] => {
    const deleteMail = db.prepare(
        `DELETE FROM raw_mails WHERE ${whereClause}`
    ).bind(...params);
    if (!isMailFlagsEnabled(env)) return [deleteMail];

    const deleteFlags = db.prepare(
        `DELETE FROM mail_flags WHERE mail_id IN (`
        + `SELECT id FROM raw_mails WHERE ${whereClause})`
    ).bind(...params);
    return [deleteFlags, deleteMail];
};

export const deleteRawMails = async (
    db: D1Database,
    env: Bindings,
    whereClause: string,
    params: (string | number)[],
): Promise<D1Result> => {
    const results = await db.batch(prepareRawMailDeleteStatements(db, env, whereClause, params));
    return results[results.length - 1];
};

export const cleanupOrphanMailFlags = async (
    db: D1Database,
    env: Bindings,
    limit = 1000,
): Promise<number> => {
    if (!isMailFlagsEnabled(env)) return 0;

    const result = await db.prepare(
        `DELETE FROM mail_flags WHERE (mail_id, flag) IN (`
        + `SELECT mf.mail_id, mf.flag FROM mail_flags mf`
        + ` LEFT JOIN raw_mails rm ON rm.id = mf.mail_id`
        + ` WHERE rm.id IS NULL LIMIT ?)`
    ).bind(limit).run();
    return result.meta.changes ?? 0;
};