import type { KTLogLeaf, KTLogStore } from './store.js';
import type { AddressIndexEntry } from './index-tree.js';
import type { SignedTreeHead } from './sth.js';
import { compareAddresses } from './index-tree.js';
import { constantTimeEqual } from './util.js';

/**
 * In-memory KTLogStore for testing and embedded servers.
 *
 * Maintains the same append-only invariants as a persistent store —
 * `appendLeaf` cannot mutate prior entries, only push.
 */
export class MemoryKTLogStore implements KTLogStore {
  private leaves: KTLogLeaf[] = [];
  private indexByAddress = new Map<string, AddressIndexEntry>();
  private sthsByTreeSize = new Map<number, SignedTreeHead[]>();
  private latestSth: SignedTreeHead | null = null;

  async appendLeaf(input: Omit<KTLogLeaf, 'index'>): Promise<number> {
    const index = this.leaves.length;
    this.leaves.push({
      ...input,
      index,
      leafHash: new Uint8Array(input.leafHash),
      bundleHash: new Uint8Array(input.bundleHash),
    });
    return index;
  }

  async getLeaves(fromIndex: number, toIndex: number): Promise<KTLogLeaf[]> {
    return this.leaves.slice(fromIndex, toIndex).map((l) => ({
      ...l,
      leafHash: new Uint8Array(l.leafHash),
      bundleHash: new Uint8Array(l.bundleHash),
    }));
  }

  async getLeaf(index: number): Promise<KTLogLeaf | null> {
    const l = this.leaves[index];
    if (!l) return null;
    return {
      ...l,
      leafHash: new Uint8Array(l.leafHash),
      bundleHash: new Uint8Array(l.bundleHash),
    };
  }

  async size(): Promise<number> {
    return this.leaves.length;
  }

  async upsertIndexEntry(entry: AddressIndexEntry): Promise<void> {
    this.indexByAddress.set(entry.address, {
      ...entry,
      bundleHash: new Uint8Array(entry.bundleHash),
    });
  }

  async tombstoneIndexEntry(address: string, latestLeafIndex: number): Promise<void> {
    const e = this.indexByAddress.get(address);
    if (!e) return;
    this.indexByAddress.set(address, {
      ...e,
      deleted: true,
      latestLeafIndex,
      bundleHash: new Uint8Array(0),
    });
  }

  async getAllIndexEntries(): Promise<AddressIndexEntry[]> {
    const all = Array.from(this.indexByAddress.values()).map((e) => ({
      ...e,
      bundleHash: new Uint8Array(e.bundleHash),
    }));
    all.sort((a, b) => compareAddresses(a.address, b.address));
    return all;
  }

  async getIndexEntry(address: string): Promise<AddressIndexEntry | null> {
    const e = this.indexByAddress.get(address);
    if (!e) return null;
    return { ...e, bundleHash: new Uint8Array(e.bundleHash) };
  }

  async saveSTH(sth: SignedTreeHead): Promise<void> {
    const cloned: SignedTreeHead = {
      treeSize: sth.treeSize,
      timestampMs: sth.timestampMs,
      rootHash: new Uint8Array(sth.rootHash),
      indexRoot: new Uint8Array(sth.indexRoot),
      logId: new Uint8Array(sth.logId),
      signature: new Uint8Array(sth.signature),
    };
    const list = this.sthsByTreeSize.get(sth.treeSize) ?? [];
    // De-duplicate (same root_hash + signature == same STH)
    const dup = list.find(
      (existing) =>
        existing.timestampMs === cloned.timestampMs &&
        constantTimeEqual(existing.rootHash, cloned.rootHash) &&
        constantTimeEqual(existing.signature, cloned.signature),
    );
    if (!dup) list.push(cloned);
    this.sthsByTreeSize.set(sth.treeSize, list);
    if (
      !this.latestSth ||
      cloned.treeSize > this.latestSth.treeSize ||
      (cloned.treeSize === this.latestSth.treeSize && cloned.timestampMs > this.latestSth.timestampMs)
    ) {
      this.latestSth = cloned;
    }
  }

  async getLatestSTH(): Promise<SignedTreeHead | null> {
    return this.latestSth ? cloneSth(this.latestSth) : null;
  }

  async getSTHByTreeSize(treeSize: number): Promise<SignedTreeHead | null> {
    const list = this.sthsByTreeSize.get(treeSize);
    if (!list || list.length === 0) return null;
    // Pick the most recent one for this tree size.
    let best = list[0]!;
    for (const s of list) if (s.timestampMs > best.timestampMs) best = s;
    return cloneSth(best);
  }

  async listSTHs(fromTimestampMs?: number, toTimestampMs?: number): Promise<SignedTreeHead[]> {
    const all: SignedTreeHead[] = [];
    for (const list of this.sthsByTreeSize.values()) {
      for (const s of list) {
        if (fromTimestampMs !== undefined && s.timestampMs < fromTimestampMs) continue;
        if (toTimestampMs !== undefined && s.timestampMs > toTimestampMs) continue;
        all.push(cloneSth(s));
      }
    }
    all.sort((a, b) => a.timestampMs - b.timestampMs);
    return all;
  }
}

function cloneSth(sth: SignedTreeHead): SignedTreeHead {
  return {
    treeSize: sth.treeSize,
    timestampMs: sth.timestampMs,
    rootHash: new Uint8Array(sth.rootHash),
    indexRoot: new Uint8Array(sth.indexRoot),
    logId: new Uint8Array(sth.logId),
    signature: new Uint8Array(sth.signature),
  };
}
