import { DatabaseSync } from "node:sqlite";
import { setTimeout as sleep } from "node:timers/promises";
import * as l from "@cerulea/lexicon";
import { Hydrant } from "./hydrant.ts";
import {
  type AtUri,
  atUri,
  Follow,
  parseAtUri,
  Star,
  starUrl,
} from "./records.ts";

export type IndexedStar = { uri: string; author: string };
const PAGE_SIZE = 25;

const EventId = l.object({ id: l.integerWith({ minimum: 0 }) });
const RecordEvent = l.object({
  type: l.literal("record"),
  record: l.object({ did: l.string, collection: l.string, rkey: l.string }),
});

// Yields a socket's text messages until it closes or the signal aborts.
// Messages queue while the consumer is busy, bounded so a flood can't
// exhaust memory.
async function* messages(
  socket: WebSocket,
  signal: AbortSignal,
): AsyncGenerator<string> {
  const queue: string[] = [];
  let closed = false;
  let wake = () => {};
  const stop = () => {
    closed = true;
    socket.close();
    wake();
  };
  socket.addEventListener("message", (event) => {
    if (typeof event.data !== "string" || queue.length >= 4096) return stop();
    queue.push(event.data);
    wake();
  });
  socket.addEventListener("error", stop);
  socket.addEventListener("close", stop);
  signal.addEventListener("abort", stop, { once: true });
  try {
    while (!closed || queue.length) {
      const next = queue.shift();
      if (next !== undefined) yield next;
      else await new Promise<void>((resolve) => wake = resolve);
    }
  } finally {
    signal.removeEventListener("abort", stop);
    stop();
  }
}

export class Index {
  private readonly viewers = new Set<string>();
  private readonly pendingViewers = new Map<string, Promise<boolean>>();
  connected = false;
  ready = false;

  constructor(readonly db: DatabaseSync, readonly hydrant: Hydrant) {
    db.exec(`
      CREATE TABLE IF NOT EXISTS stars (uri TEXT PRIMARY KEY, author TEXT NOT NULL, tid TEXT NOT NULL);
      CREATE INDEX IF NOT EXISTS stars_order ON stars (tid DESC, uri DESC);
      CREATE TABLE IF NOT EXISTS follows (uri TEXT PRIMARY KEY, viewer TEXT NOT NULL, subject TEXT NOT NULL);
      CREATE INDEX IF NOT EXISTS follows_viewer ON follows (viewer, subject);
      CREATE TABLE IF NOT EXISTS sync_state (key TEXT PRIMARY KEY, value TEXT NOT NULL);
    `);
  }

  private transaction<T>(body: () => T): T {
    this.db.exec("BEGIN IMMEDIATE");
    try {
      const result = body();
      this.db.exec("COMMIT");
      return result;
    } catch (error) {
      this.db.exec("ROLLBACK");
      throw error;
    }
  }

  // The last stream event applied.
  private get cursor(): number {
    const row = this.db.prepare(
      "SELECT value FROM sync_state WHERE key = 'cursor'",
    ).get();
    return Number(row?.value ?? -1);
  }
  private set cursor(id: number) {
    this.db.prepare("INSERT OR REPLACE INTO sync_state VALUES ('cursor', ?)")
      .run(String(id));
  }

  // Mirrors a record's current state; a missing or invalid one is removed.
  private mirror(uri: string, value: unknown): void {
    const collection = parseAtUri(uri)?.collection;
    if (collection === Star.id) this.star(uri, value);
    else if (collection === Follow.id) this.follow(uri, value);
  }

  star(uri: string, value: unknown): void {
    const parsed = Star.validate(value);
    const at = parseAtUri(uri);
    if (
      parsed.success && at?.collection === Star.id && starUrl(parsed.value.url)
    ) {
      this.db.prepare("INSERT OR REPLACE INTO stars VALUES (?, ?, ?)").run(
        uri,
        at.did,
        at.rkey,
      );
    } else this.removeStar(uri);
  }
  removeStar(uri: string): void {
    this.db.prepare("DELETE FROM stars WHERE uri = ?").run(uri);
  }

  follow(uri: string, value: unknown): void {
    const parsed = Follow.validate(value);
    const at = parseAtUri(uri);
    if (
      parsed.success && at?.collection === Follow.id &&
      parsed.value.subject !== at.did
    ) {
      this.db.prepare("INSERT OR REPLACE INTO follows VALUES (?, ?, ?)").run(
        uri,
        at.did,
        parsed.value.subject,
      );
    } else this.removeFollow(uri);
  }
  removeFollow(uri: string): void {
    this.db.prepare("DELETE FROM follows WHERE uri = ?").run(uri);
  }

  following(viewer: string, subject: string): string[] {
    return this.db.prepare(
      "SELECT uri FROM follows WHERE viewer = ? AND subject = ?",
    ).all(viewer, subject).map((row) => String(row.uri));
  }

  // Fetches one row past the page to learn whether there is a next one.
  page(
    viewer: string,
    before?: string,
  ): { stars: IndexedStar[]; next?: string } {
    const rows = this.db.prepare(`SELECT uri, author FROM stars WHERE
      (author = ? OR author IN (SELECT subject FROM follows WHERE viewer = ?))
      AND (? IS NULL OR (tid, uri) < (?, ?))
      ORDER BY tid DESC, uri DESC LIMIT ${PAGE_SIZE + 1}`).all(
      viewer,
      viewer,
      before ?? null,
      before?.split("/").at(-1) ?? "",
      before ?? "",
    ).map((row) => ({ uri: String(row.uri), author: String(row.author) }));
    const stars = rows.slice(0, PAGE_SIZE);
    return {
      stars,
      next: rows.length > PAGE_SIZE ? stars.at(-1)?.uri : undefined,
    };
  }

  // Whether the viewer's follows are known, preparing them on first sight.
  ensureViewer(did: string): Promise<boolean> {
    if (this.viewers.has(did)) return Promise.resolve(true);
    const pending = this.pendingViewers.get(did);
    if (pending) return pending;
    const task = (async () => {
      try {
        await this.hydrant.track(did);
        await this.reconcileFollows(did);
        this.viewers.add(did);
        return true;
      } catch (error) {
        console.error("Could not prepare viewer", did, error);
        return false;
      } finally {
        this.pendingViewers.delete(did);
      }
    })();
    this.pendingViewers.set(did, task);
    return task;
  }

  async reconcileFollows(viewer: string): Promise<void> {
    const follows = await Array.fromAsync(
      this.hydrant.allRecords(viewer, Follow.id),
    );
    this.transaction(() => {
      this.db.prepare("DELETE FROM follows WHERE viewer = ?").run(viewer);
      for (const { uri, value } of follows) this.follow(uri, value);
    });
  }

  async bootstrap(): Promise<void> {
    for await (const did of this.hydrant.repos()) {
      for await (
        const { uri, value } of this.hydrant.allRecords(did, Star.id)
      ) {
        this.star(uri, value);
      }
    }
    const sessions = this.db.prepare(
      "SELECT DISTINCT did FROM browser_session WHERE expires > ?",
    ).all(Date.now());
    for (const row of sessions) await this.ensureViewer(String(row.did));
    this.ready = true;
  }

  // Every star, and the follows of people who use this site.
  private relevant({ did, collection }: AtUri): boolean {
    return collection === Star.id ||
      (collection === Follow.id && this.viewers.has(did));
  }

  async apply(raw: unknown): Promise<void> {
    const event = EventId.validate(raw);
    if (!event.success) throw new Error("Invalid Hydrant event");
    const { id } = event.value;
    if (id <= this.cursor) return;
    const change = RecordEvent.validate(raw);
    const uri = change.success && this.relevant(change.value.record)
      ? atUri(change.value.record)
      : undefined;
    const current = uri ? await this.hydrant.record(uri) : undefined;
    this.transaction(() => {
      if (uri) this.mirror(uri, current?.value);
      this.cursor = id;
    });
  }

  async run(signal: AbortSignal): Promise<void> {
    while (!signal.aborted) {
      try {
        if (!this.ready) await this.bootstrap();
        await this.consume(this.hydrant.stream(this.cursor + 1), signal);
      } catch (error) {
        if (!signal.aborted) console.error("Hydrant sync failed", error);
      }
      if (!signal.aborted) await sleep(1000);
    }
  }

  private async consume(url: URL, signal: AbortSignal): Promise<void> {
    const socket = new WebSocket(url);
    const deadline = setTimeout(() => socket.close(), 15000);
    socket.addEventListener("open", () => {
      clearTimeout(deadline);
      this.connected = true;
    }, { once: true });
    try {
      for await (const message of messages(socket, signal)) {
        await this.apply(JSON.parse(message));
      }
    } finally {
      clearTimeout(deadline);
      this.connected = false;
    }
  }
}