import postgres, { type Sql } from 'postgres';
import type { PrekeyStore } from '@shade/server';
import { toBase64, fromBase64 } from '@shade/core';
import { ensurePrekeyServerTables } from './ensure-tables.js';

/**
 * PostgreSQL-backed PrekeyStore for the Shade Prekey Server.
 *
 * Concurrent-safe one-time prekey consumption using FOR UPDATE SKIP LOCKED,
 * so multiple server instances can share the same PostgreSQL.
 */
export class PostgresPrekeyStore implements PrekeyStore {
  private constructor(
    private readonly sql: Sql,
    private readonly ownsConnection: boolean,
  ) {}

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

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

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

  async saveIdentity(address: string, identitySigningKey: Uint8Array, identityDHKey: Uint8Array): Promise<void> {
    await this.sql`
      INSERT INTO shade_server_identities (address, identity_signing_key, identity_dh_key, last_activity_at)
      VALUES (${address}, ${toBase64(identitySigningKey)}, ${toBase64(identityDHKey)}, ${Date.now()})
      ON CONFLICT (address) DO UPDATE SET
        identity_signing_key = EXCLUDED.identity_signing_key,
        identity_dh_key = EXCLUDED.identity_dh_key,
        last_activity_at = EXCLUDED.last_activity_at
    `;
  }

  async getIdentity(address: string): Promise<{ identitySigningKey: Uint8Array; identityDHKey: Uint8Array } | null> {
    const rows = await this.sql<Array<{ identity_signing_key: string; identity_dh_key: string }>>`
      SELECT identity_signing_key, identity_dh_key FROM shade_server_identities WHERE address = ${address}
    `;
    if (rows.length === 0) return null;
    return {
      identitySigningKey: fromBase64(rows[0]!.identity_signing_key),
      identityDHKey: fromBase64(rows[0]!.identity_dh_key),
    };
  }

  async saveSignedPreKey(address: string, keyId: number, publicKey: Uint8Array, signature: Uint8Array): Promise<void> {
    await this.sql`
      INSERT INTO shade_server_signed_prekeys (address, key_id, public_key, signature)
      VALUES (${address}, ${keyId}, ${toBase64(publicKey)}, ${toBase64(signature)})
      ON CONFLICT (address) DO UPDATE SET
        key_id = EXCLUDED.key_id,
        public_key = EXCLUDED.public_key,
        signature = EXCLUDED.signature
    `;
  }

  async getSignedPreKey(address: string): Promise<{ keyId: number; publicKey: Uint8Array; signature: Uint8Array } | null> {
    const rows = await this.sql<Array<{ key_id: number; public_key: string; signature: string }>>`
      SELECT key_id, public_key, signature FROM shade_server_signed_prekeys WHERE address = ${address}
    `;
    if (rows.length === 0) return null;
    return {
      keyId: rows[0]!.key_id,
      publicKey: fromBase64(rows[0]!.public_key),
      signature: fromBase64(rows[0]!.signature),
    };
  }

  async saveOneTimePreKeys(address: string, keys: Array<{ keyId: number; publicKey: Uint8Array }>): Promise<void> {
    if (keys.length === 0) return;
    await this.sql.begin(async (sql) => {
      for (const k of keys) {
        await sql`
          INSERT INTO shade_server_one_time_prekeys (address, key_id, public_key)
          VALUES (${address}, ${k.keyId}, ${toBase64(k.publicKey)})
        `;
      }
    });
  }

  async consumeOneTimePreKey(address: string): Promise<{ keyId: number; publicKey: Uint8Array } | null> {
    // Atomic consume with FOR UPDATE SKIP LOCKED for concurrent safety
    const rows = await this.sql<Array<{ key_id: number; public_key: string }>>`
      DELETE FROM shade_server_one_time_prekeys
      WHERE id = (
        SELECT id FROM shade_server_one_time_prekeys
        WHERE address = ${address}
        ORDER BY id
        LIMIT 1
        FOR UPDATE SKIP LOCKED
      )
      RETURNING key_id, public_key
    `;
    if (rows.length === 0) return null;
    return {
      keyId: rows[0]!.key_id,
      publicKey: fromBase64(rows[0]!.public_key),
    };
  }

  async getOneTimePreKeyCount(address: string): Promise<number> {
    const rows = await this.sql<Array<{ count: string }>>`
      SELECT COUNT(*)::text as count FROM shade_server_one_time_prekeys WHERE address = ${address}
    `;
    return parseInt(rows[0]!.count, 10);
  }

  async deleteAll(address: string): Promise<void> {
    await this.sql.begin(async (sql) => {
      await sql`DELETE FROM shade_server_identities WHERE address = ${address}`;
      await sql`DELETE FROM shade_server_signed_prekeys WHERE address = ${address}`;
      await sql`DELETE FROM shade_server_one_time_prekeys WHERE address = ${address}`;
    });
  }

  // ─── Stale cleanup ──────────────────────────────────────

  async touchIdentity(address: string): Promise<void> {
    await this.sql`
      UPDATE shade_server_identities
      SET last_activity_at = ${Date.now()}
      WHERE address = ${address}
    `;
  }

  async purgeStaleIdentities(olderThanMs: number): Promise<number> {
    const cutoff = Date.now() - olderThanMs;
    const rows = await this.sql<Array<{ address: string }>>`
      SELECT address FROM shade_server_identities
      WHERE last_activity_at < ${cutoff}
    `;
    if (rows.length === 0) return 0;

    const addresses = rows.map((r) => r.address);
    await this.sql.begin(async (sql) => {
      await sql`DELETE FROM shade_server_identities WHERE address = ANY(${addresses})`;
      await sql`DELETE FROM shade_server_signed_prekeys WHERE address = ANY(${addresses})`;
      await sql`DELETE FROM shade_server_one_time_prekeys WHERE address = ANY(${addresses})`;
    });
    return addresses.length;
  }
}
