import * as Y from "yjs";
import { Awareness, encodeAwarenessUpdate, applyAwarenessUpdate } from "y-protocols/awareness";

export const DOC_COLLECTION = "com.jakelazaroff.y-pds.doc";
export const UPDATE_COLLECTION = "com.jakelazaroff.y-pds.update";
export const AWARENESS_COLLECTION = "com.jakelazaroff.y-pds.awareness";
export const WEBRTC_COLLECTION = "com.jakelazaroff.y-pds.webrtc";

const encode = (/** @type {Uint8Array} */ update) => btoa(String.fromCharCode(...update));
const decode = (/** @type {string} */ b64) => Uint8Array.from(atob(b64), c => c.charCodeAt(0));

const TID_CHARS = "234567abcdefghijklmnopqrstuvwxyz";

const s32encode = (/** @type {number} */ i) => {
  let s = "";
  while (i) {
    s = TID_CHARS[i % 32] + s;
    i = Math.floor(i / 32);
  }
  return s;
};

/** Generate an AT Protocol TID: a 13-character, lexicographically sortable record key. */
export const genTid = () => {
  const timestamp = Date.now() * 1000 + Math.floor(Math.random() * 1000);
  const clockId = Math.floor(Math.random() * 1024);
  return s32encode(timestamp).padStart(11, "2") + s32encode(clockId).padStart(2, "2");
};

const MSG_DOC = 0;
const MSG_AWARENESS = 1;

const ICE_SERVERS = [
  { urls: "stun:stun.l.google.com:19302" },
];

export class YPdsProvider {
  /** @type {Y.Doc} */
  #ydoc;

  /** @type {Map<string, string>} */
  #pds = new Map();

  /** @param {string} did */
  async #getPDS(did) {
    let pds = this.#pds.get(did);
    if (pds) return pds;

    const url = did.startsWith("did:web:")
      ? `https://${did.slice("did:web:".length)}/.well-known/did.json`
      : `https://plc.directory/${did}`;
    const doc = await fetch(url).then(r => r.json());

    pds = doc.service?.find(s => s.type === "AtprotoPersonalDataServer")?.serviceEndpoint;
    if (!pds) throw new Error(`no PDS found for ${did}`);

    this.#pds.set(did, pds);
    return pds;
  }

  /**
   * @params {object} options
   * @params {string} [options.rkey]
   * @params {string} options.repo
   * @params {string} options.collection
   */
  async #getRecord({ repo, collection, rkey }) {
    const pds = await this.#getPDS(repo);
    const params = new URLSearchParams({ repo, collection, rkey });
    const res = await fetch(`${pds}/xrpc/com.atproto.repo.getRecord?${params}`);
    const data = await res.json();
    if (!res.ok) throw new Error(data.message);
    return data;
  }

  /**
   * @params {object} options
   * @params {string} options.repo
   * @params {string} options.collection
   * @params {string} [options.cursor]
   */
  async #listRecords({ repo, collection, cursor }) {
    const pds = await this.#getPDS(repo);
    const params = new URLSearchParams({ repo, collection, limit: "100" });
    if (cursor) params.set("cursor", cursor);
    const res = await fetch(`${pds}/xrpc/com.atproto.repo.listRecords?${params}`);
    const data = await res.json();
    if (!res.ok) throw new Error(data.message);
    return data;
  }

  /**
   * @params {object} options
   * @params {string} options.repo
   * @params {string} options.collection
   * @params {object} options.record
   * @params {string} [options.rkey]
   */
  async #upsertRecord({ repo, collection, rkey, record }) {
    const pds = await this.#getPDS(repo);
    const verb = rkey ? "com.atproto.repo.putRecord" : "com.atproto.repo.createRecord";
    const res = await fetch(`${pds}/xrpc/${verb}`, {
      method: "POST",
      headers: {
        "content-type": "application/json",
        "x-atsw-did": repo,
      },
      body: JSON.stringify({ repo, collection, rkey, record: { $type: collection, ...record } }),
    });
    const data = await res.json();
    if (!res.ok) throw new Error(data.message);
    return data;
  }

  #ownerDid = "";
  #rkey = "";

  /** @type {Array<{uri: string, cid: string, repo: string, value: {docId: string, update: string, createdAt: string}}>} */
  #allRecords = [];

  /** @type {Map<string, {repo: string, uri: string, cid: string, value: object}>} */
  #awarenessMap = new Map();

  /** @type {Set<WebSocket>} */
  #sockets = new Set();

  // Subscribe to multiple Jetstream instances at once: any single instance can
  // silently drop or lag commits from some repos (e.g. jetstream2.us-east
  // missing us-west repos), so we fan out and dedupe. Yjs updates are
  // idempotent, but deduping keeps side effects (record list, load(), WebRTC
  // signaling) from firing once per instance.
  #jetstreams = [
    "wss://jetstream1.us-east.bsky.network/subscribe",
    "wss://jetstream2.us-east.bsky.network/subscribe",
    "wss://jetstream1.us-west.bsky.network/subscribe",
  ];

  /** @type {Map<string, number>} per-endpoint cursor (time_us of last event) */
  #cursors = new Map();

  /** @type {Set<string>} commit keys already processed, for cross-instance dedupe */
  #seen = new Set();
  /** @type {string[]} insertion order for bounding #seen */
  #seenOrder = [];

  /** @type {{ editors: string[], [key: string]: unknown } | null} */
  #doc = null;

  #destroyed = false;

  /**
   * @param {object} options
   * @param {Y.Doc} options.ydoc
   * @param {string} options.pds
   * @param {string} options.atUri
   * @param {string} options.did
   * @param {Awareness} [options.awareness]
   * @param {string} [options.jetstream]
   */
  /** @type {Map<string, { pc: RTCPeerConnection }>} */
  #peers = new Map();

  /** @type {((editors: string[]) => void) | null} */
  onMembersChange = null;

  /** @type {(() => void) | null} */
  onRecordAdded = null;

  /** @type {(() => void) | null} */
  onAwarenessChange = null;

  /** @type {(() => void) | null} */
  onPauseChange = null;

  /** @type {((did: string) => void) | null} */
  onPeerConnect = null;

  /** @type {((did: string) => void) | null} */
  onPeerDisconnect = null;

  #paused = false;
  get paused() { return this.#paused; }

  #throttle = 5000;

  /** @type {Uint8Array[]} */
  #pendingUpdates = [];
  #updateTimerId = null;

  #awarenessDirty = false;
  #awarenessTimerId = null;

  constructor({ ydoc, pds, atUri, did, awareness, jetstream, jetstreams, throttle }) {
    this.#pds.set(did, pds);
    this.#ydoc = ydoc;

    const [ownerDid, , rkey] = atUri.slice("at://".length).split("/");
    this.#ownerDid = ownerDid;
    this.#rkey = rkey;
    this.did = did;
    this.awareness = awareness ?? new Awareness(ydoc);
    if (jetstreams) this.#jetstreams = jetstreams;
    else if (jetstream) this.#jetstreams = [jetstream];
    if (throttle !== undefined) this.#throttle = throttle;

    this.awareness.on("change", this.#onAwarenessChange);
    this.#ydoc.on("update", this.#onUpdate);
  }

  async load() {
    try {
      const data = await this.#getRecord({
        repo: this.#ownerDid,
        collection: DOC_COLLECTION,
        rkey: this.#rkey,
      });
      this.#doc = data.value;
      this.onMembersChange?.(this.#doc.editors);
    } catch {
      // if the fetch fails and the document is not in our own repo,
      // throw an error rather than attempting to create it
      if (this.did !== this.#ownerDid) throw new Error("Document not found");

      // create the document
      this.#doc = { docId: this.#rkey, editors: [], createdAt: new Date().toISOString() };
      await this.#upsertRecord({
        repo: this.did,
        collection: DOC_COLLECTION,
        rkey: this.#rkey,
        record: this.#doc,
      });
      this.onMembersChange?.(this.#doc.editors);
    }

    // fetch updates from the owner and all editors
    const repos = [this.#ownerDid].concat(this.#doc.editors);
    const updates = await Promise.all(repos.map(repo => this.#fetchUpdates(repo)));

    // apply all the updates
    Y.transact(this.#ydoc, () => {
      for (const { value } of updates.flat()) {
        Y.applyUpdate(this.#ydoc, decode(value.update));
      }
    }, this);

    this.#allRecords = updates.flat();

    // Persist a consolidated copy in our own PDS so we keep the full document
    // even if other editors later delete their records.
    await this.#consolidate();

    this.#subscribe();
  }

  async #consolidate() {
    // Nothing to back up if every record is already ours.
    if (!this.#allRecords.some(r => r.repo !== this.did)) return;

    const merged = Y.encodeStateAsUpdate(this.#ydoc);
    await this.#upsertRecord({
      repo: this.did,
      collection: UPDATE_COLLECTION,
      // Fixed rkey = the doc's own key, so putRecord overwrites rather than
      // accumulating a new record on every load.
      rkey: this.#rkey,
      record: { docId: this.#rkey, update: encode(merged), createdAt: new Date().toISOString() },
    });
  }

  async #fetchUpdates(repo) {
    const records = [];
    let cursor;

    do {
      const res = await this.#listRecords({ repo, collection: UPDATE_COLLECTION, cursor });
      records.push(
        ...res.records
          .filter(r => r.value.docId === this.#rkey)
          .map(r => ({ ...r, repo })),
      );
      cursor = res.cursor;
    } while (cursor);

    return records;
  }

  /**
   * @param {Uint8Array} update
   * @param {unknown} origin
   */
  #flushUpdate = async () => {
    this.#updateTimerId = null;
    if (this.#pendingUpdates.length === 0) return;
    const merged = this.#pendingUpdates.length === 1
      ? this.#pendingUpdates[0]
      : Y.mergeUpdates(this.#pendingUpdates);
    this.#pendingUpdates = [];
    const createdAt = new Date().toISOString();
    const data = await this.#upsertRecord({
      repo: this.did,
      collection: UPDATE_COLLECTION,
      record: { docId: this.#rkey, update: encode(merged), createdAt },
    });
    this.#allRecords.push({
      uri: data.uri,
      cid: data.cid,
      repo: this.did,
      value: { docId: this.#rkey, update: encode(merged), createdAt },
    });
    this.onRecordAdded?.();
  };

  #onUpdate = (update, origin) => {
    if (origin === this) return;
    this.#pendingUpdates.push(update);
    if (!this.#paused && this.#updateTimerId == null) {
      this.#updateTimerId = setTimeout(this.#flushUpdate, this.#throttle);
    }
  };

  /** @param {{ added: number[], updated: number[] }} changes */
  #flushAwareness = async () => {
    this.#awarenessTimerId = null;
    if (!this.#awarenessDirty) return;
    this.#awarenessDirty = false;
    const update = encodeAwarenessUpdate(this.awareness, [this.awareness.clientID]);
    const createdAt = new Date().toISOString();
    const data = await this.#upsertRecord({
      repo: this.did,
      collection: AWARENESS_COLLECTION,
      rkey: this.#rkey,
      record: { docId: this.#rkey, awareness: encode(update), createdAt },
    });
    this.#awarenessMap.set(this.did, {
      repo: this.did,
      uri: data.uri,
      cid: data.cid,
      value: { docId: this.#rkey, awareness: encode(update), createdAt },
    });
    this.onAwarenessChange?.();
  };

  #onAwarenessChange = ({ added, updated }) => {
    const ours = this.awareness.clientID;
    if (!added.includes(ours) && !updated.includes(ours)) return;
    this.#awarenessDirty = true;
    if (!this.#paused && this.#awarenessTimerId == null) {
      this.#awarenessTimerId = setTimeout(this.#flushAwareness, this.#throttle);
    }
  };

  getMembers() {
    return this.#doc.editors;
  }

  getAwarenessRecords() {
    return [...this.#awarenessMap.values()];
  }

  getDebugInfo() {
    return {
      doc: this.#doc,
      docUri: `at://${this.#ownerDid}/${DOC_COLLECTION}/${this.#rkey}`,
      ownerDid: this.#ownerDid,
      rkey: this.#rkey,
      records: [...this.#allRecords],
    };
  }

  pause() {
    if (this.#paused) return;
    this.#paused = true;
    clearTimeout(this.#updateTimerId);
    this.#updateTimerId = null;
    clearTimeout(this.#awarenessTimerId);
    this.#awarenessTimerId = null;
    this.onPauseChange?.();
  }

  async play() {
    if (!this.#paused) return;
    this.#paused = false;

    // Re-fetch all repos and apply only records we haven't seen yet,
    // merged into one update so the ydoc sees a single transaction.
    const repos = [this.#ownerDid, ...this.#doc.editors];
    const fetched = (await Promise.all(repos.map(r => this.#fetchUpdates(r)))).flat();
    const seenUris = new Set(this.#allRecords.map(r => r.uri));
    const newRecords = fetched.filter(r => !seenUris.has(r.uri));
    if (newRecords.length > 0) {
      const merged = Y.mergeUpdates(newRecords.map(r => decode(r.value.update)));
      Y.applyUpdate(this.#ydoc, merged, this);
      this.#allRecords.push(...newRecords);
      this.onRecordAdded?.();
    }

    // Flush buffered local edits as one record.
    await this.#flushUpdate();

    this.onPauseChange?.();
  }

  getPeers() {
    return [...this.#peers.keys()];
  }

  async connect(dids) {
    await Promise.all(dids.map(async did => {
      if (did === this.did || this.#peers.has(did)) return;

      const pc = new RTCPeerConnection({ iceServers: ICE_SERVERS });
      const channel = pc.createDataChannel("yjs");
      this.#peers.set(did, { pc });
      this.#setupDataChannel(did, channel);

      const icePromise = this.#waitForIce(pc);
      const offer = await pc.createOffer();
      await pc.setLocalDescription(offer);
      await icePromise;

      await this.#upsertRecord({
        repo: this.did,
        collection: WEBRTC_COLLECTION,
        record: {
          target: did,
          docId: this.#rkey,
          type: "offer",
          sdp: pc.localDescription.sdp,
          createdAt: new Date().toISOString(),
        },
      });
    }));
  }

  /** Wait for ICE gathering to finish. Call before setLocalDescription to avoid missing candidates. */
  #waitForIce(pc) {
    return new Promise(resolve => {
      if (pc.iceGatheringState === "complete") { resolve(); return; }
      const timer = setTimeout(resolve, 10_000);
      pc.addEventListener("icecandidate", function check(e) {
        if (!e.candidate) {
          clearTimeout(timer);
          pc.removeEventListener("icecandidate", check);
          resolve();
        }
      });
    });
  }

  /** @param {string} did @param {RTCDataChannel} channel */
  #setupDataChannel(did, channel) {
    channel.binaryType = "arraybuffer";

    const send = (type, data) => {
      if (channel.readyState !== "open") return;
      const msg = new Uint8Array(1 + data.byteLength);
      msg[0] = type;
      msg.set(data, 1);
      channel.send(msg);
    };

    channel.onopen = () => {
      this.onPeerConnect?.(did);
      send(MSG_DOC, Y.encodeStateAsUpdate(this.#ydoc));
      send(MSG_AWARENESS, encodeAwarenessUpdate(this.awareness, [this.awareness.clientID]));
    };

    channel.onmessage = e => {
      const msg = new Uint8Array(e.data);
      const payload = msg.slice(1);
      if (msg[0] === MSG_DOC) Y.applyUpdate(this.#ydoc, payload, "webrtc");
      else if (msg[0] === MSG_AWARENESS) applyAwarenessUpdate(this.awareness, payload, "webrtc");
    };

    const onUpdate = (update, origin) => {
      if (origin === "webrtc" || origin === this) return;
      send(MSG_DOC, update);
    };
    this.#ydoc.on("update", onUpdate);

    const onAwarenessChange = ({ added, updated }) => {
      const own = [...added, ...updated].filter(id => id === this.awareness.clientID);
      if (own.length > 0) send(MSG_AWARENESS, encodeAwarenessUpdate(this.awareness, own));
    };
    this.awareness.on("change", onAwarenessChange);

    channel.onclose = () => {
      this.#ydoc.off("update", onUpdate);
      this.awareness.off("change", onAwarenessChange);
      this.#peers.delete(did);
      this.onPeerDisconnect?.(did);
    };
  }

  async #handleOffer(did, record) {
    if (this.#peers.has(did)) return;

    const pc = new RTCPeerConnection({ iceServers: ICE_SERVERS });
    this.#peers.set(did, { pc });
    pc.ondatachannel = e => this.#setupDataChannel(did, e.channel);

    const icePromise = this.#waitForIce(pc);
    await pc.setRemoteDescription({ type: "offer", sdp: record.sdp });
    const answer = await pc.createAnswer();
    await pc.setLocalDescription(answer);
    await icePromise;

    await this.#upsertRecord({
      repo: this.did,
      collection: WEBRTC_COLLECTION,
      record: {
        target: did,
        docId: this.#rkey,
        type: "answer",
        sdp: pc.localDescription.sdp,
        createdAt: new Date().toISOString(),
      },
    });
  }

  async #handleAnswer(did, record) {
    const peer = this.#peers.get(did);
    if (!peer || peer.pc.signalingState !== "have-local-offer") return;
    await peer.pc.setRemoteDescription({ type: "answer", sdp: record.sdp });
  }

  async setMembers(editors) {
    this.#doc = { ...this.#doc, editors };
    await this.#upsertRecord({
      repo: this.did,
      collection: DOC_COLLECTION,
      rkey: this.#rkey,
      record: this.#doc,
    });
  }

  #subscribe() {
    this.#closeSockets();
    for (const endpoint of this.#jetstreams) this.#connect(endpoint);
  }

  /** Close every socket without triggering its reconnect (intentional teardown). */
  #closeSockets() {
    for (const ws of this.#sockets) {
      ws.onclose = null;
      ws.close();
    }
    this.#sockets.clear();
  }

  /** Record a commit key; returns false if it was already seen on another instance. */
  #once(key) {
    if (this.#seen.has(key)) return false;
    this.#seen.add(key);
    this.#seenOrder.push(key);
    if (this.#seenOrder.length > 2000) this.#seen.delete(this.#seenOrder.shift());
    return true;
  }

  #connect(endpoint) {
    const url = new URL(endpoint);
    url.searchParams.append("wantedCollections", UPDATE_COLLECTION);
    url.searchParams.append("wantedCollections", AWARENESS_COLLECTION);
    url.searchParams.append("wantedCollections", DOC_COLLECTION);
    url.searchParams.append("wantedCollections", WEBRTC_COLLECTION);
    url.searchParams.append("wantedDids", this.#ownerDid);
    for (const editor of this.#doc.editors) url.searchParams.append("wantedDids", editor);

    const cursor = this.#cursors.get(endpoint);
    if (cursor) url.searchParams.set("cursor", String(cursor));

    const ws = new WebSocket(url);
    this.#sockets.add(ws);
    ws.onerror = () => {};
    ws.onclose = () => {
      this.#sockets.delete(ws);
      if (this.#destroyed) return;
      setTimeout(() => {
        if (!this.#destroyed) this.#connect(endpoint);
      }, 1000);
    };

    ws.onmessage = e => {
      const event = JSON.parse(e.data);
      this.#cursors.set(endpoint, event.time_us);
      if (event.kind !== "commit") return;

      // The same commit arrives from every instance; process it only once.
      const c = event.commit;
      if (!this.#once(`${event.did}/${c.collection}/${c.rkey}/${c.cid ?? c.operation}`)) return;

      this.#handleCommit(event);
    };
  }

  async #handleCommit(event) {
      switch (event.commit.collection) {
        case DOC_COLLECTION: {
          if (event.did !== this.#ownerDid || event.commit.rkey !== this.#rkey) return;
          await this.load();
          break;
        }

        case UPDATE_COLLECTION: {
          if (event.commit.record.docId !== this.#rkey) return;
          if (event.did !== this.#ownerDid && !this.#doc.editors.includes(event.did)) return;

          if (event.commit.operation !== "create") return;

          if (this.#paused) break;

          // Own writes are already tracked optimistically in #flushUpdate
          if (event.did !== this.did) {
            this.#allRecords.push({
              uri: `at://${event.did}/${UPDATE_COLLECTION}/${event.commit.rkey}`,
              cid: event.commit.cid,
              repo: event.did,
              value: event.commit.record,
            });
            this.onRecordAdded?.();
          }
          Y.applyUpdate(this.#ydoc, decode(event.commit.record.update), this);
          break;
        }

        case AWARENESS_COLLECTION: {
          if (event.commit.record.docId !== this.#rkey) return;
          if (event.did !== this.#ownerDid && !this.#doc.editors.includes(event.did)) return;

          if (this.#paused) break;

          this.#awarenessMap.set(event.did, {
            repo: event.did,
            uri: `at://${event.did}/${AWARENESS_COLLECTION}/${event.commit.rkey}`,
            cid: event.commit.cid,
            value: event.commit.record,
          });
          this.onAwarenessChange?.();
          if (event.did === this.did) break;
          applyAwarenessUpdate(this.awareness, decode(event.commit.record.awareness), "remote");
          break;
        }

        case WEBRTC_COLLECTION: {
          const rec = event.commit.record;
          if (rec.target !== this.did) return;
          if (rec.docId !== this.#rkey) return;
          if (event.did !== this.#ownerDid && !this.#doc.editors.includes(event.did)) return;
          if (Date.now() - new Date(rec.createdAt).getTime() > 120_000) return;

          if (rec.type === "offer") this.#handleOffer(event.did, rec).catch(console.error);
          else if (rec.type === "answer") this.#handleAnswer(event.did, rec).catch(console.error);
          break;
        }
      }
  }

  destroy() {
    this.#destroyed = true;
    this.#ydoc.off("update", this.#onUpdate);
    this.awareness.off("change", this.#onAwarenessChange);
    clearTimeout(this.#updateTimerId);
    clearTimeout(this.#awarenessTimerId);
    this.#flushUpdate();
    this.awareness.destroy();
    this.#closeSockets();
    for (const { pc } of this.#peers.values()) pc.close();
    this.#peers.clear();
  }
}