import { hashKey, type QueryClient } from "@tanstack/react-query";

import type { CacheScopeSnapshot } from "../../../contracts/server-state-scope.ts";
import { OPTIMISTIC_LAYER_BOUNDS } from "../../../contracts/server-state.ts";

export type OptimisticLayerLease = Readonly<{
  commit(): void;
  rollback(): void;
  markUncertain(): void;
  reconcile(resolution: "APPLIED" | "NOT_APPLIED"): void;
}>;

type Layer = {
  id: number;
  status: "pending" | "uncertain" | "committed";
  apply(value: unknown): unknown;
};

type OptimisticLayerScope = Pick<CacheScopeSnapshot, "isCurrent"> &
  Partial<Pick<CacheScopeSnapshot, "signal">>;

type EntryState = {
  queryKey: readonly unknown[];
  scope: OptimisticLayerScope;
  base: unknown;
  layers: Layer[];
  disposeScopeListener: (() => void) | null;
};

export function createOptimisticLayerRuntime(queryClient: QueryClient) {
  const entries = new Map<string, EntryState>();
  let nextId = 1;
  let writing = false;

  function removeEntry(key: string, entry: EntryState): void {
    if (entries.get(key) === entry) entries.delete(key);
    entry.disposeScopeListener?.();
    entry.disposeScopeListener = null;
  }

  function writeProjection(entry: EntryState, value: unknown): void {
    writing = true;
    try {
      queryClient.setQueryData(entry.queryKey, value);
    } finally {
      writing = false;
    }
  }

  queryClient.getQueryCache().subscribe((event) => {
    if (
      writing ||
      event.type !== "updated" ||
      !entries.has(event.query.queryHash)
    ) {
      return;
    }
    const entry = entries.get(event.query.queryHash);
    if (!entry) return;
    entry.base = event.query.state.data;
    project(event.query.queryHash, entry);
  });

  function project(key: string, entry: EntryState): void {
    if (!entry.scope.isCurrent()) {
      removeEntry(key, entry);
      queryClient.removeQueries({ queryKey: entry.queryKey, exact: true });
      return;
    }
    if (entry.base === undefined && entry.layers.length === 0) {
      removeEntry(key, entry);
      queryClient.removeQueries({ queryKey: entry.queryKey, exact: true });
      return;
    }
    let value = entry.base;
    try {
      for (const layer of entry.layers) value = layer.apply(value);
    } catch {
      removeEntry(key, entry);
      return;
    }
    writeProjection(entry, value);
  }

  function collapse(key: string, entry: EntryState): void {
    while (entry.layers[0]?.status === "committed") {
      const committed = entry.layers.shift();
      if (!committed) break;
      entry.base = committed.apply(entry.base);
    }
    project(key, entry);
    if (entry.layers.length === 0) removeEntry(key, entry);
  }

  return Object.freeze({
    begin<Input>(
      queryKey: readonly unknown[],
      input: Input,
      update: (previous: unknown, input: Input) => unknown,
      scope: OptimisticLayerScope,
    ): OptimisticLayerLease | null {
      if (!scope.isCurrent()) return null;
      const current = queryClient.getQueryData(queryKey);
      const key = hashKey(queryKey);
      let entry = entries.get(key);
      if (!entry) {
        entry = {
          queryKey,
          scope,
          base: current,
          layers: [],
          disposeScopeListener: null,
        };
        entries.set(key, entry);
        if (scope.signal) {
          const selectedEntry = entry;
          const discardScope = () => {
            removeEntry(key, selectedEntry);
            queryClient.removeQueries({ queryKey, exact: true });
          };
          scope.signal.addEventListener("abort", discardScope, { once: true });
          entry.disposeScopeListener = () =>
            scope.signal?.removeEventListener("abort", discardScope);
          if (scope.signal.aborted || !scope.isCurrent()) {
            discardScope();
            return null;
          }
        }
      } else if (entry.scope !== scope) {
        return null;
      }
      // §11.5. Overflow falls back to pessimistic execution; an existing
      // layer is never silently evicted to make room for a new one.
      if (entry.layers.length >= OPTIMISTIC_LAYER_BOUNDS.maxLayersPerQueryKey) {
        if (entry.layers.length === 0) removeEntry(key, entry);
        return null;
      }
      const layer: Layer = {
        id: nextId,
        status: "pending",
        apply: (value) => update(value, input),
      };
      let projected: unknown;
      try {
        projected = entry.base;
        for (const existingLayer of entry.layers) {
          projected = existingLayer.apply(projected);
        }
        projected = layer.apply(projected);
      } catch (error) {
        if (entry.layers.length > 0) return null;
        removeEntry(key, entry);
        throw error;
      }
      if (
        estimateLayerBytes(projected) >
        OPTIMISTIC_LAYER_BOUNDS.maxSingleLayerBytes
      ) {
        if (entry.layers.length === 0) removeEntry(key, entry);
        return null;
      }
      nextId += 1;
      entry.layers.push(layer);
      writeProjection(entry, projected);
      let state: "pending" | "uncertain" | "settled" = "pending";
      const selectedLayer = () => {
        if (entries.get(key) !== entry) return undefined;
        return entry.layers.find((candidate) => candidate.id === layer.id);
      };
      return Object.freeze({
        commit() {
          if (state !== "pending") return;
          state = "settled";
          const selected = selectedLayer();
          if (!selected) return;
          selected.status = "committed";
          collapse(key, entry);
        },
        rollback() {
          if (state !== "pending") return;
          state = "settled";
          if (!selectedLayer()) return;
          entry.layers = entry.layers.filter(
            (candidate) => candidate.id !== layer.id,
          );
          collapse(key, entry);
        },
        markUncertain() {
          if (state !== "pending") return;
          state = "uncertain";
          const selected = selectedLayer();
          if (!selected) return;
          selected.status = "uncertain";
          collapse(key, entry);
        },
        reconcile(resolution) {
          if (state !== "uncertain") return;
          state = "settled";
          const selected = selectedLayer();
          if (!selected) return;
          if (resolution === "APPLIED") {
            selected.status = "committed";
          } else {
            entry.layers = entry.layers.filter(
              (candidate) => candidate.id !== layer.id,
            );
          }
          collapse(key, entry);
        },
      });
    },
  });
}

const encoder = new TextEncoder();

/**
 * A local, bounded estimate of one optimistic projection. This is not the
 * §10.4 query result measurement: it only decides whether a rollback snapshot
 * stays inside the layer budget, and it stops as soon as the budget is passed.
 */
function estimateLayerBytes(value: unknown): number {
  let total = 0;
  const stack: unknown[] = [value];
  let visited = 0;
  while (stack.length > 0) {
    if (visited++ > 4_096) return Number.POSITIVE_INFINITY;
    if (total > OPTIMISTIC_LAYER_BOUNDS.maxSingleLayerBytes) return total;
    const current = stack.pop();
    if (typeof current === "string") {
      total += encoder.encode(current).byteLength;
    } else if (typeof current === "number" || typeof current === "boolean") {
      total += 8;
    } else if (Array.isArray(current)) {
      total += 8;
      for (const item of current) stack.push(item);
    } else if (current && typeof current === "object") {
      total += 8;
      for (const [key, item] of Object.entries(current)) {
        total += encoder.encode(key).byteLength;
        stack.push(item);
      }
    }
  }
  return total;
}