import postgres, { type Sql } from 'postgres';
import { toBase64, fromBase64 } from '@shade/core';
import {
  type AddressIndexEntry,
  type KTLogLeaf,
  type KTLogStore,
  type SignedTreeHead,
  compareAddresses,
} from '@shade/key-transparency';
import { ensureKTLogTables } from './ensure-tables.js';

/**
 * PostgreSQL-backed KTLogStore.
 *
 * Append-only invariant is enforced two ways:
 *  - Application code only ever runs INSERT against `shade_kt_leaves`.
 *  - The table has a CHECK constraint and a trigger that reject UPDATE /
 *    DELETE in production deployments. (See `ensureKTLogTables`.)
 *
 * Multiple server instances may share the same KT log; row-level
 * `SELECT … FOR UPDATE` on the count row prevents two writers from
 * picking the same leaf index.
 */
export class PostgresKTLogStore implements KTLogStore {
  private constructor(
    private readonly sql: Sql,
    private readonly ownsConnection: boolean,
  ) {}

  static async create(connectionString: string): Promise<PostgresKTLogStore> {
    const sql = postgres(connectionString);
    const store = new PostgresKTLogStore(sql, true);
    await ensureKTLogTables(sql);
    return store;
  }

  static async fromClient(sql: Sql): Promise<PostgresKTLogStore> {
    const store = new PostgresKTLogStore(sql, false);
    await ensureKTLogTables(sql);
    return store;
  }

  async close(): Promise<void> {
    if (this.ownsConnection) await this.sql.end();
  }

  async appendLeaf(input: Omit<KTLogLeaf, 'index'>): Promise<number> {
    return this.sql.begin(async (sql) => {
      const rows = await sql<Array<{ next_index: string }>>`
        SELECT COALESCE(MAX(leaf_index), -1)::bigint + 1 AS next_index
        FROM shade_kt_leaves
        FOR UPDATE
      `;
      const idx = parseInt(rows[0]!.next_index, 10);
      await sql`
        INSERT INTO shade_kt_leaves
          (leaf_index, leaf_hash, timestamp_ms, operation, address, bundle_hash)
        VALUES
          (${idx}, ${toBase64(input.leafHash)}, ${input.timestampMs}, ${input.operation},
           ${input.address}, ${toBase64(input.bundleHash)})
      `;
      return idx;
    });
  }

  async getLeaves(fromIndex: number, toIndex: number): Promise<KTLogLeaf[]> {
    if (toIndex <= fromIndex) return [];
    const rows = await this.sql<
      Array<{
        leaf_index: string;
        leaf_hash: string;
        timestamp_ms: string;
        operation: number;
        address: string;
        bundle_hash: string;
      }>
    >`
      SELECT leaf_index, leaf_hash, timestamp_ms, operation, address, bundle_hash
      FROM shade_kt_leaves
      WHERE leaf_index >= ${fromIndex} AND leaf_index < ${toIndex}
      ORDER BY leaf_index ASC
    `;
    return rows.map((r) => ({
      index: parseInt(r.leaf_index, 10),
      leafHash: fromBase64(r.leaf_hash),
      timestampMs: parseInt(r.timestamp_ms, 10),
      operation: r.operation,
      address: r.address,
      bundleHash: fromBase64(r.bundle_hash),
    }));
  }

  async getLeaf(index: number): Promise<KTLogLeaf | null> {
    const rows = await this.sql<
      Array<{
        leaf_index: string;
        leaf_hash: string;
        timestamp_ms: string;
        operation: number;
        address: string;
        bundle_hash: string;
      }>
    >`
      SELECT leaf_index, leaf_hash, timestamp_ms, operation, address, bundle_hash
      FROM shade_kt_leaves
      WHERE leaf_index = ${index}
    `;
    if (rows.length === 0) return null;
    const r = rows[0]!;
    return {
      index: parseInt(r.leaf_index, 10),
      leafHash: fromBase64(r.leaf_hash),
      timestampMs: parseInt(r.timestamp_ms, 10),
      operation: r.operation,
      address: r.address,
      bundleHash: fromBase64(r.bundle_hash),
    };
  }

  async size(): Promise<number> {
    const rows = await this.sql<Array<{ count: string }>>`
      SELECT COUNT(*)::text AS count FROM shade_kt_leaves
    `;
    return parseInt(rows[0]!.count, 10);
  }

  async upsertIndexEntry(entry: AddressIndexEntry): Promise<void> {
    await this.sql`
      INSERT INTO shade_kt_index (address, latest_leaf_index, bundle_hash, deleted)
      VALUES (${entry.address}, ${entry.latestLeafIndex}, ${toBase64(entry.bundleHash)}, ${entry.deleted})
      ON CONFLICT (address) DO UPDATE SET
        latest_leaf_index = EXCLUDED.latest_leaf_index,
        bundle_hash = EXCLUDED.bundle_hash,
        deleted = EXCLUDED.deleted
    `;
  }

  async tombstoneIndexEntry(address: string, latestLeafIndex: number): Promise<void> {
    await this.sql`
      UPDATE shade_kt_index
      SET deleted = TRUE,
          latest_leaf_index = ${latestLeafIndex},
          bundle_hash = ''
      WHERE address = ${address}
    `;
  }

  async getAllIndexEntries(): Promise<AddressIndexEntry[]> {
    const rows = await this.sql<
      Array<{
        address: string;
        latest_leaf_index: string;
        bundle_hash: string;
        deleted: boolean;
      }>
    >`
      SELECT address, latest_leaf_index, bundle_hash, deleted
      FROM shade_kt_index
      ORDER BY address ASC
    `;
    const entries = rows.map((r) => ({
      address: r.address,
      latestLeafIndex: parseInt(r.latest_leaf_index, 10),
      bundleHash: r.bundle_hash ? fromBase64(r.bundle_hash) : new Uint8Array(0),
      deleted: r.deleted,
    }));
    // Postgres' ORDER BY may use locale-aware collation in some configs;
    // re-sort using our explicit byte-lex compare to match the in-memory
    // canonical ordering.
    entries.sort((a, b) => compareAddresses(a.address, b.address));
    return entries;
  }

  async getIndexEntry(address: string): Promise<AddressIndexEntry | null> {
    const rows = await this.sql<
      Array<{
        latest_leaf_index: string;
        bundle_hash: string;
        deleted: boolean;
      }>
    >`
      SELECT latest_leaf_index, bundle_hash, deleted
      FROM shade_kt_index WHERE address = ${address}
    `;
    if (rows.length === 0) return null;
    const r = rows[0]!;
    return {
      address,
      latestLeafIndex: parseInt(r.latest_leaf_index, 10),
      bundleHash: r.bundle_hash ? fromBase64(r.bundle_hash) : new Uint8Array(0),
      deleted: r.deleted,
    };
  }

  async saveSTH(sth: SignedTreeHead): Promise<void> {
    await this.sql`
      INSERT INTO shade_kt_sths
        (tree_size, timestamp_ms, root_hash, index_root, log_id, signature)
      VALUES
        (${sth.treeSize}, ${sth.timestampMs}, ${toBase64(sth.rootHash)},
         ${toBase64(sth.indexRoot)}, ${toBase64(sth.logId)}, ${toBase64(sth.signature)})
      ON CONFLICT (tree_size, timestamp_ms, signature) DO NOTHING
    `;
  }

  async getLatestSTH(): Promise<SignedTreeHead | null> {
    const rows = await this.sql<
      Array<{
        tree_size: string;
        timestamp_ms: string;
        root_hash: string;
        index_root: string;
        log_id: string;
        signature: string;
      }>
    >`
      SELECT tree_size, timestamp_ms, root_hash, index_root, log_id, signature
      FROM shade_kt_sths
      ORDER BY tree_size DESC, timestamp_ms DESC
      LIMIT 1
    `;
    if (rows.length === 0) return null;
    return rowToSth(rows[0]!);
  }

  async getSTHByTreeSize(treeSize: number): Promise<SignedTreeHead | null> {
    const rows = await this.sql<
      Array<{
        tree_size: string;
        timestamp_ms: string;
        root_hash: string;
        index_root: string;
        log_id: string;
        signature: string;
      }>
    >`
      SELECT tree_size, timestamp_ms, root_hash, index_root, log_id, signature
      FROM shade_kt_sths
      WHERE tree_size = ${treeSize}
      ORDER BY timestamp_ms DESC
      LIMIT 1
    `;
    if (rows.length === 0) return null;
    return rowToSth(rows[0]!);
  }

  async listSTHs(fromTimestampMs?: number, toTimestampMs?: number): Promise<SignedTreeHead[]> {
    const from = fromTimestampMs ?? 0;
    const to = toTimestampMs ?? Number.MAX_SAFE_INTEGER;
    const rows = await this.sql<
      Array<{
        tree_size: string;
        timestamp_ms: string;
        root_hash: string;
        index_root: string;
        log_id: string;
        signature: string;
      }>
    >`
      SELECT tree_size, timestamp_ms, root_hash, index_root, log_id, signature
      FROM shade_kt_sths
      WHERE timestamp_ms >= ${from} AND timestamp_ms <= ${to}
      ORDER BY timestamp_ms ASC
    `;
    return rows.map(rowToSth);
  }
}

function rowToSth(r: {
  tree_size: string;
  timestamp_ms: string;
  root_hash: string;
  index_root: string;
  log_id: string;
  signature: string;
}): SignedTreeHead {
  return {
    treeSize: parseInt(r.tree_size, 10),
    timestampMs: parseInt(r.timestamp_ms, 10),
    rootHash: fromBase64(r.root_hash),
    indexRoot: fromBase64(r.index_root),
    logId: fromBase64(r.log_id),
    signature: fromBase64(r.signature),
  };
}
