import type { Bindings, JetstreamEvent } from "../types";
import { resolvePds } from "../utils/resolver";
import { ingestDocument, deleteDocument } from "../utils/ingest";

const JETSTREAM_URL =
	"wss://jetstream2.us-east.bsky.network/subscribe?wantedCollections=site.standard.document";
const ALARM_INTERVAL_MS = 30_000;
const CURSOR_SAVE_INTERVAL_MS = 10_000;
const CURSOR_SAVE_MESSAGE_COUNT = 100;

export class JetstreamConsumer implements DurableObject {
	private state: DurableObjectState;
	private env: Bindings;
	private ws: WebSocket | null = null;
	private cursor: string | null = null;
	private messageCount = 0;
	private lastCursorSave = 0;
	private pdsCache: Map<string, string | null> = new Map();

	constructor(state: DurableObjectState, env: Bindings) {
		this.state = state;
		this.env = env;
	}

	async fetch(request: Request): Promise<Response> {
		const url = new URL(request.url);

		switch (url.pathname) {
			case "/start":
				return this.handleStart();
			case "/status":
				return this.handleStatus();
			case "/stop":
				return this.handleStop();
			default:
				return new Response("Not found", { status: 404 });
		}
	}

	async alarm(): Promise<void> {
		if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
			console.log("Alarm: WebSocket not connected, reconnecting...");
			await this.connect();
		}
		// Reschedule alarm
		await this.state.storage.setAlarm(Date.now() + ALARM_INTERVAL_MS);
	}

	private async handleStart(): Promise<Response> {
		if (this.ws && this.ws.readyState === WebSocket.OPEN) {
			return Response.json({
				status: "already_connected",
				cursor: this.cursor,
			});
		}

		await this.connect();
		await this.state.storage.setAlarm(Date.now() + ALARM_INTERVAL_MS);

		return Response.json({ status: "started", cursor: this.cursor });
	}

	private handleStatus(): Response {
		return Response.json({
			connected: this.ws?.readyState === WebSocket.OPEN,
			cursor: this.cursor,
			messageCount: this.messageCount,
		});
	}

	private async handleStop(): Promise<Response> {
		if (this.ws) {
			this.ws.close();
			this.ws = null;
		}
		await this.state.storage.deleteAlarm();
		return Response.json({ status: "stopped" });
	}

	private async connect(): Promise<void> {
		// Close existing connection if any
		if (this.ws) {
			try {
				this.ws.close();
			} catch {}
			this.ws = null;
		}

		// Load cursor from storage
		if (!this.cursor) {
			this.cursor =
				(await this.state.storage.get<string>("cursor")) || null;
		}

		const url = this.cursor
			? `${JETSTREAM_URL}&cursor=${this.cursor}`
			: JETSTREAM_URL;

		console.log(`Connecting to Jetstream: ${url}`);

		try {
			const ws = new WebSocket(url);

			ws.addEventListener("message", (event) => {
				this.onMessage(event.data as string);
			});

			ws.addEventListener("close", () => {
				console.log("Jetstream WebSocket closed");
				this.ws = null;
				// Alarm will handle reconnection
			});

			ws.addEventListener("error", (event) => {
				console.error("Jetstream WebSocket error:", event);
				try {
					ws.close();
				} catch {}
				this.ws = null;
			});

			this.ws = ws;
		} catch (error) {
			console.error("Failed to connect to Jetstream:", error);
			this.ws = null;
		}
	}

	private async onMessage(data: string): Promise<void> {
		try {
			const event = JSON.parse(data) as JetstreamEvent;

			// Update cursor from every event
			this.cursor = String(event.time_us);

			// Only process commit events
			if (event.kind !== "commit") return;

			const { commit } = event;
			if (commit.collection !== "site.standard.document") return;

			// PDS filter: skip bridgy noise
			const pds = await this.resolvePdsWithCache(event.did);
			if (!pds || pds.includes("brid.gy")) return;

			if (
				commit.operation === "create" ||
				commit.operation === "update"
			) {
				await ingestDocument(this.env.DB, this.env.RESOLUTION_QUEUE, {
					did: event.did,
					rkey: commit.rkey,
					collection: commit.collection,
					cid: commit.cid,
					record: commit.record,
				});
			} else if (commit.operation === "delete") {
				await deleteDocument(this.env.DB, {
					did: event.did,
					collection: commit.collection,
					rkey: commit.rkey,
				});
			}

			this.messageCount++;

			// Periodically save cursor
			await this.maybeSaveCursor();
		} catch (error) {
			console.error("Error processing Jetstream message:", error);
		}
	}

	private async resolvePdsWithCache(did: string): Promise<string | null> {
		// Fast in-memory cache (cleared on DO eviction)
		if (this.pdsCache.has(did)) {
			return this.pdsCache.get(did)!;
		}

		const pds = await resolvePds(this.env.DB, did);
		this.pdsCache.set(did, pds);

		// Keep in-memory cache bounded
		if (this.pdsCache.size > 10_000) {
			const firstKey = this.pdsCache.keys().next().value;
			if (firstKey) this.pdsCache.delete(firstKey);
		}

		return pds;
	}

	private async maybeSaveCursor(): Promise<void> {
		if (!this.cursor) return;

		const now = Date.now();
		const shouldSave =
			this.messageCount % CURSOR_SAVE_MESSAGE_COUNT === 0 ||
			now - this.lastCursorSave > CURSOR_SAVE_INTERVAL_MS;

		if (shouldSave) {
			await this.state.storage.put("cursor", this.cursor);
			this.lastCursorSave = now;
		}
	}
}