import { type Selectable, type Transaction } from "kysely";
import type { DB } from "../database/out/db";
import { log } from "../util";
import { sleep } from "bun";
import { db } from "../database/db";
import { AssertError, Value } from "@sinclair/typebox/value";
import type { TSchema, Static } from "@sinclair/typebox";
import { RateLimitError, secToMs } from "@blade-and-brawn/domain";
import { EventsService, type EventSource } from "./events";

type RetryType = "error" | "ratelimit";

type EventProcessor<P extends TSchema, S extends TSchema> =
    (payload: Static<P>, setState: (state: Static<S>) => Promise<void>, id: string) => Promise<void>;

export type EventTypeOptions<P extends TSchema, S extends TSchema> = {
    schemas: { payload: P, state: S },
    processor: EventProcessor<P, S>,
    dedupeKey: (payload: Static<P>) => string | undefined
};

type EventQueueCfg = {
    error: {
        limit: number,
        initialRetryDelayMs: number,
    },
    concurrency: {
        limit: number,
    },
};

class AlreadyClaimedEventError extends Error { }
class ConcurrencyLimitReachedError extends Error { }

export function defineEventType<P extends TSchema, S extends TSchema>(opt: EventTypeOptions<P, S>): EventTypeOptions<P, S> {
    return opt;
}

function addJitter(delay: number, jitterFactor: number): number {
    const jitterMultiplier = 1 + (Math.random() - 0.5) * jitterFactor;
    return delay * jitterMultiplier;
}

function exponentialDelay(initialDelay: number, tries: number, cap: number = Infinity) {
    return Math.min(initialDelay * (2 ** tries), cap);
}

export class EventQueue<G extends string, T extends Record<string, EventTypeOptions<any, any>>> {
    private static readonly DEQUEUE_RETRY_DELAY_MS = secToMs(10);
    private static readonly DEQUEUE_RETRY_LIMIT = 3;
    private static readonly RATE_LIMIT_RETRY_LIMIT = 15;
    private static readonly JITTER_FACTOR = 0.2;

    readonly cfg: EventQueueCfg;
    readonly group: G;
    readonly types: T;

    private readonly Events = new EventsService();

    private _lastCleanDate: Date = new Date(0);
    get lastCleanDate() { return this._lastCleanDate; };

    constructor(opt: { group: G, cfg: EventQueueCfg, types: T }) {
        this.cfg = opt.cfg;
        this.group = opt.group;
        this.types = opt.types;
    }

    async drain(): Promise<void> {
        log.trace("draining queue");
        let consecutiveFailures = 0;
        while (true) {
            try {
                const event = await this.dequeue();
                if (!event) break;
                consecutiveFailures = 0;
            }
            catch (err) {
                if (err instanceof AlreadyClaimedEventError) {
                    log.warn("event processing already claimed by another worker, skipping");
                    continue;
                }
                if (err instanceof ConcurrencyLimitReachedError) {
                    await sleep(EventQueue.DEQUEUE_RETRY_DELAY_MS);
                    continue;
                }

                if (++consecutiveFailures <= EventQueue.DEQUEUE_RETRY_LIMIT) {
                    log.warn({ err, consecutiveFailures }, "dequeue failed, retrying")
                    await sleep(EventQueue.DEQUEUE_RETRY_DELAY_MS);
                    continue;
                }

                log.error({ err, consecutiveFailures }, "failed to drain queue");
                throw err;
            }
        };
    }

    async dequeue() {
        log.trace("dequeueing");

        const front = await this.front();
        if (!front) return;

        const result = await db.updateTable("events")
            .set({ started_at: new Date(), heartbeat_at: new Date() })
            .where("id", "=", front.id)
            .where("started_at", "is", null)
            .where((eb) =>
                eb(
                    eb.selectFrom("events")
                        .select(eb.fn.countAll().as("count"))
                        .where("group", "=", this.group)
                        .where("started_at", "is not", null)
                        .where("ended_at", "is", null),
                    "<",
                    this.cfg.concurrency.limit
                )
            )
            .executeTakeFirstOrThrow();

        // Guard against another worker/machine having already claimed this event and guard against max concurrency
        if (result.numUpdatedRows === 0n) {
            const frontFresh = await db.selectFrom("events")
                .select("started_at")
                .where("id", "=", front.id)
                .executeTakeFirst();

            if (!frontFresh) throw new AlreadyClaimedEventError();
            if (frontFresh?.started_at) throw new AlreadyClaimedEventError();
            throw new ConcurrencyLimitReachedError();
        };

        log.trace({ name: front.id }, "dequeuing event");
        await this.process(front);

        return front;
    }

    async enqueue<K extends keyof T & string>({ source, type, payload }: {
        source: EventSource,
        type: K,
        payload: Static<T[K]["schemas"]["payload"]>
    }): Promise<void> {
        if (!(type in this.types)) throw new Error(`event type not registered in this queue (${this.group}): ${type}`);
        const cfg = this.types[type as keyof T]!;

        const dedupeKey = cfg.dedupeKey(payload);
        log.info({ payload, dedupeKey }, "enqueuing");

        await db.insertInto("events")
            .values({
                group: this.group,
                type,
                source,
                payload,
                dedupe_key: dedupeKey ?? null,
            })
            .onConflict((oc) => oc
                .columns(["group", "type", "dedupe_key"])
                .where("dedupe_key", "is not", null)
                .where("ended_at", "is", null)
                .doNothing()
            )
            .returning(["id"])
            .executeTakeFirst();
    }

    async front(trx?: Transaction<DB>) {
        const result = await (trx ?? db).selectFrom("events")
            .select(["id", "type", "payload", "available_at", "errors", "rate_limits"])
            .where("group", "=", this.group)
            .where("started_at", "is", null)
            .where("available_at", "<=", new Date())
            .orderBy("created_at", "asc")
            .limit(1)
            .executeTakeFirst();
        return result;
    }

    async clean(opt: { filter?: { type?: string } } = {}) {
        try {
            await this.Events.clean({ filter: { group: this.group, ...opt.filter } });
        }
        finally {
            this._lastCleanDate = new Date();
        }
    }

    private async process(event: Pick<Selectable<DB["events"]>, "id" | "type" | "errors" | "payload" | "rate_limits">) {
        log.info({ id: event.id, type: event.type, payload: event.payload }, "processing event");
        if (!(event.type in this.types)) throw new Error(`event type not registered in this queue (${this.group}): ${event.type}`);
        const typeCfg = this.types[event.type as keyof T]!;

        try {
            Value.Assert(typeCfg.schemas.payload, event.payload);
            await typeCfg.processor(event.payload, (s) => this.Events.setState(event.id, s), event.id);
            await this.Events.resolve(event.id);
        }
        catch (err) {
            await this.Events.setLastError(event.id, err);

            if (err instanceof RateLimitError && event.rate_limits < EventQueue.RATE_LIMIT_RETRY_LIMIT) {
                log.warn(
                    { rateLimits: event.rate_limits, delay: err.retryDelayMs, source: err.source, payload: err.payload },
                    "encountered a rate limit error"
                );
                await this.scheduleRetry(event.id, err.retryDelayMs, "ratelimit");
                return;
            }

            if (err instanceof AssertError) {
                log.error({ errors: event.errors }, "could not process event: malformed payload");
                await this.Events.fail(event.id);
                return;
            }

            if (event.errors < this.cfg.error.limit) {
                const retryDelayMs = addJitter(
                    exponentialDelay(this.cfg.error.initialRetryDelayMs, event.errors),
                    EventQueue.JITTER_FACTOR
                );
                log.warn({ err, errors: event.errors, retryDelayMs: retryDelayMs }, "encountered an error during event processing, retrying")
                await this.scheduleRetry(event.id, retryDelayMs, "error");
                return;
            }

            log.error({ err, errors: event.errors, rateLimits: event.rate_limits }, "event failed");
            await this.Events.fail(event.id);
        }
    }

    private async scheduleRetry(id: string, delayMs: number, type: RetryType) {
        await db.updateTable("events")
            .set((eb) => ({
                errors: type === "error" ? eb("errors", "+", 1) : eb.ref("errors"),
                rate_limits: type === "ratelimit" ? eb("rate_limits", "+", 1) : eb.ref("rate_limits"),
                heartbeat_at: null,
                started_at: null,
                ended_at: null,
                has_failed: false,
                available_at: new Date(Date.now() + delayMs)
            }))
            .where("id", "=", id)
            .execute();
    }
}