import { Plugin } from "@opencode/plugin";
import { QuotaTracker } from "./quota.js";
import { QuotaRpc } from "./quota-rpc.js";

export default Plugin.define({
    id: "stanp-opencode.quota",
    async setup(context) {
        const controller = new AbortController();
        const quota = new QuotaTracker(context, async () => {
            await registration.events.emit("updated", {});
        });
        const registration = await context.rpc.register(QuotaRpc, {
            status: async () => quota.status() ?? null,
        });
        await quota.init();
        const updates = (async () => {
            for await (const event of context.event.subscribe({ signal: controller.signal })) {
                if (
                    event.location &&
                    (event.location.directory !== context.location.directory ||
                        ("workspaceID" in event.location
                            ? event.location.workspaceID
                            : undefined) !== context.location.workspaceID)
                )
                    continue;
                if (
                    event.type === "session.execution.succeeded" ||
                    event.type === "session.execution.failed" ||
                    event.type === "session.execution.interrupted" ||
                    event.type === "session.idle"
                ) {
                    await quota.complete(event.data.sessionID);
                } else if (event.type === "credential.updated") {
                    await quota.authChanged();
                } else if (
                    event.type === "credential.switched" &&
                    String(event.data.integrationID) === "openai"
                ) {
                    await quota.authChanged(true);
                }
            }
        })().catch((error) => {
            if (!controller.signal.aborted)
                console.error("stanp-opencode quota update failed", error);
        });
        await context.session.hook("model.request", async (event) => quota.start(event.sessionID));
        return async () => {
            controller.abort();
            await updates;
            await quota.dispose();
            await registration.dispose();
        };
    },
});