import WebSocket from "ws";

export type JetstreamEvent = {
  did: string;
  time_us: number;
} & (CommitEvent | AccountEvent | IdentityEvent);

type CommitEvent = {
  kind: "commit";
  commit:
    | {
        operation: "create" | "update";
        record: unknown;
        rev: string;
        collection: string;
        rkey: string;
        cid: string;
      }
    | {
        operation: "delete";
        rev: string;
        collection: string;
        rkey: string;
      };
};

type IdentityEvent = {
  kind: "identity";
  identity: {
    did: string;
    handle: string;
    seq: number;
    time: string;
  };
};

export type AccountEvent = {
  kind: "account";
  account: {
    active: boolean;
    did: string;
    seq: number;
    time: string;
    status?: "takendown" | "suspended" | "deleted" | "deactivated";
  };
};

export interface JetstreamConfig {
  instanceUrl: string;
  wantedCollections: string[];
  cursor?: number;
  onEvent: (evt: JetstreamEvent) => Promise<void>;
  onError: (err: unknown) => void;
  setCursor?: (cursorUs: number) => Promise<void>;
}

export class Jetstream {
  private ws?: WebSocket;
  private config: JetstreamConfig;
  private isStarted = false;
  private isDestroyed = false;
  private cursor?: number;
  private lastCursorWrite = 0;

  constructor(config: JetstreamConfig) {
    this.config = config;
    this.cursor = config.cursor;
  }

  private buildUrl(): string {
    const params = new URLSearchParams();
    for (const collection of this.config.wantedCollections) {
      params.append("wantedCollections", collection);
    }
    if (this.cursor !== undefined) {
      params.append("cursor", this.cursor.toString());
    }
    return `${this.config.instanceUrl}?${params.toString()}`;
  }

  start(): void {
    if (this.isStarted) return;
    this.isStarted = true;
    this.isDestroyed = false;

    const url = this.buildUrl();
    console.log(`Connecting to Jetstream: ${url}`);
    this.ws = new WebSocket(url);

    this.ws.on("open", () => {
      console.log("Jetstream connection opened");
    });

    this.ws.on("message", async (data) => {
      try {
        const event: JetstreamEvent = JSON.parse(data.toString());

        // Update cursor (throttled to every 30s)
        if (event.time_us !== undefined && this.config.setCursor) {
          this.cursor = event.time_us;
          const now = Date.now();
          if (now - this.lastCursorWrite >= 30000) {
            this.lastCursorWrite = now;
            console.log(`Saving cursor: ${event.time_us}`);
            await this.config.setCursor(event.time_us);
          }
        }

        await this.config.onEvent(event);
      } catch (err) {
        this.config.onError(err);
      }
    });

    this.ws.on("error", (err) => {
      this.config.onError(err);
    });

    this.ws.on("close", (code, reason) => {
      this.isStarted = false;
      if (!this.isDestroyed) {
        console.log(`Jetstream closed (code: ${code}, reason: ${reason}). Reconnecting in 5s...`);
        setTimeout(() => this.start(), 5000);
      }
    });
  }

  destroy(): void {
    this.isDestroyed = true;
    if (this.ws) {
      this.ws.close();
      this.isStarted = false;
      console.log("Jetstream destroyed");
    }
  }
}
