import { BinaryWriter, BinaryReader } from './graph/CondensedGraphBinary.js';
import { CondensedGraph } from './CondensedGraph.js';

const MAGIC = 0x41524231; // ARB1

const CONFIG_STRIP_KEYS = new Set(['_compiled', '_compileErrors', '_compileWarnings']);
const COMPILED_STRIP_KEYS = new Set(['source']);

function toSnapshotSafe(value, stripKeys) {
  const seen = new WeakSet();
  return JSON.parse(JSON.stringify(value, (key, val) => {
    if (stripKeys && stripKeys.has(key)) return undefined;
    if (typeof val === 'function') return undefined;
    if (!val || typeof val !== 'object') return val;
    if (seen.has(val)) return undefined;
    seen.add(val);
    return val;
  }));
}

function serializeArbiterSnapshot(arbiter, options = {}) {
  const writer = new BinaryWriter();
  writer.writeUint32(MAGIC);
  writer.writeUint32(1);

  const graph = options.graph || arbiter.snapshotGraph;
  if (!graph) {
    throw new Error('Snapshot graph is not available');
  }

  const graphBinary = graph.toBinary();
  writer.writeUint32(graphBinary.byteLength);
  writer.writeBytes(new Uint8Array(graphBinary));

  const relationEntries = [];
  for (const [name, config] of arbiter.relationConfigs.entries()) {
    const configPayload = config ? toSnapshotSafe(config, CONFIG_STRIP_KEYS) : null;
    const compiledPayload = config?._compiled ? toSnapshotSafe(config._compiled, COMPILED_STRIP_KEYS) : null;
    relationEntries.push({
      name,
      config: configPayload,
      compiled: compiledPayload
    });
  }

  const nodeEntries = [];
  for (const node of arbiter.nodes.values()) {
    if (!node) continue;
    nodeEntries.push({
      id: node.id,
      key: node.key,
      type: node.type || null,
      data: node.data ? toSnapshotSafe(node.data) : null
    });
  }

  const snapshot = {
    relations: relationEntries,
    relationIdSnapshot: arbiter.keyManager.getRelationIdSnapshot(),
    nodes: nodeEntries
  };

  const json = JSON.stringify(snapshot);
  const encoded = new TextEncoder().encode(json);
  writer.writeUint32(encoded.byteLength);
  writer.writeBytes(encoded);
  return writer.toArrayBuffer();
}

function deserializeArbiterSnapshot(buffer, arbiterFactory) {
  const reader = new BinaryReader(buffer);
  const magic = reader.readUint32();
  if (magic !== MAGIC) throw new Error('Invalid arbiter snapshot');
  const version = reader.readUint32();
  if (version !== 1) throw new Error(`Unsupported arbiter snapshot version ${version}`);

  const graphLength = reader.readUint32();
  const graphBytes = reader.readBytes(graphLength);
  const graphCopy = new Uint8Array(graphLength);
  graphCopy.set(graphBytes);
  const graph = CondensedGraph.fromBinary(graphCopy.buffer);

  const jsonLength = reader.readUint32();
  const jsonBytes = reader.readBytes(jsonLength);
  const json = new TextDecoder().decode(jsonBytes);
  const payload = JSON.parse(json);

  const arbiter = arbiterFactory();
  arbiter.snapshotEnabled = true;
  arbiter.snapshotGraph = graph;
  arbiter._snapshotReadOnly = true;
  arbiter.relationConfigs.clear();

  arbiter.nodes.clear();
  arbiter.nodeIdByKey.clear();
  arbiter.keyByNodeId.clear();
  let maxNodeId = -1;
  for (let nodeId = 0; nodeId < graph.nodeIds.length; nodeId++) {
    const key = graph.getNodeKey(nodeId);
    if (!key) continue;
    arbiter.nodeIdByKey.set(key, nodeId);
    arbiter.keyByNodeId.set(nodeId, key);
    // Key nodes by BOTH key and nodeId to match live engine (NodeManager uses nodeId)
    arbiter.nodes.set(nodeId, { id: nodeId, key, type: null, data: null });
    if (nodeId > maxNodeId) maxNodeId = nodeId;
  }
  if (maxNodeId >= 0) {
    arbiter.nextNodeId = maxNodeId + 1;
  }

  if (payload.relationIdSnapshot) {
    arbiter.keyManager.loadRelationIdSnapshot(payload.relationIdSnapshot);
  }

  if (Array.isArray(payload.relations)) {
    for (const entry of payload.relations) {
      if (!entry || !entry.name) continue;
      const config = entry.config || null;
      if (config && entry.compiled) {
        config._compiled = entry.compiled;
      }
      arbiter.relationConfigs.set(entry.name, config);
    }
  }

  if (Array.isArray(payload.nodes)) {
    for (const entry of payload.nodes) {
      if (!entry || entry.key === undefined || entry.key === null) continue;
      const nodeId = arbiter.nodeIdByKey.get(entry.key);
      if (nodeId === undefined) continue;
      const node = arbiter.nodes.get(nodeId) || { id: nodeId, key: entry.key };
      node.type = entry.type || null;
      node.data = entry.data === undefined ? null : entry.data;
      arbiter.nodes.set(nodeId, node);
    }
  }

  return arbiter;
}

export { serializeArbiterSnapshot, deserializeArbiterSnapshot };
