import { instrumentDrizzleClient } from "@kubiks/otel-drizzle";
import { drizzle } from "drizzle-orm/node-postgres";
import pg from "pg";
import config from "@/config";
import logger from "@/logging";
import {
  installDbErrorSafetyNet,
  withTransactionRetry,
  wrapPoolWithRetry,
} from "./retry";
import * as schema from "./schemas";
import {
  DATABASE_URL_VAULT_REF_ENV,
  getDatabaseUrlFromVault,
  isReadonlyVaultEnabled,
} from "./vault-database-url";

/** Type for database transactions */
export type Transaction = Parameters<
  Parameters<ReturnType<typeof getDb>["transaction"]>[0]
>[0];

/**
 * Initialize the database connection pool.
 * This must be called before any database operations.
 *
 * The function:
 * 1. Checks if ARCHESTRA_DATABASE_URL_VAULT_REF is set AND READONLY_VAULT is enabled
 * 2. If so, reads database URL from Vault
 * 3. Otherwise falls back to ARCHESTRA_DATABASE_URL or DATABASE_URL env vars
 * 4. Creates the connection pool with keepalive settings
 *
 * @throws Error if database URL is not configured anywhere
 */
export async function initializeDatabase(): Promise<void> {
  if (db) {
    return; // Already initialized
  }

  // Share one in-flight init between concurrent callers so only a single
  // pool is ever created. On failure the stored promise is cleared, so a
  // later call retries instead of rejecting forever.
  if (!initPromise) {
    initPromise = doInitializeDatabase().catch((error) => {
      initPromise = null;
      throw error;
    });
  }
  return initPromise;
}

/**
 * Get the database instance.
 * @throws Error if database is not initialized
 */
export function getDb() {
  if (!db) {
    throw new Error(
      "Database not initialized. Call initializeDatabase() first.",
    );
  }
  return db;
}

/**
 * Run a database transaction with retry around the whole transaction.
 *
 * Pool-level retry only applies to standalone `pool.query()` calls. Drizzle
 * transactions use a checked-out client, so transient connection failures must
 * retry the entire transaction callback.
 */
export async function withDbTransaction<T>(
  callback: (tx: Transaction) => Promise<T>,
): Promise<T> {
  return withTransactionRetry(() => getDb().transaction(callback));
}

/**
 * Check if the database connection is healthy by executing a simple query.
 * Returns true if the database is reachable, false otherwise.
 *
 * Uses a 3-second timeout to prevent hanging probes under high load.
 * This is called every 10 seconds by K8s readiness probes (per Helm config).
 *
 * Note: pool.query() is wrapped with retry logic (see {@link wrapPoolWithRetry}).
 * This is intentional for health checks — a single transient blip (e.g.
 * ECONNREFUSED during a brief network hiccup) will be retried rather than
 * immediately marking the pod as unhealthy and causing unnecessary restarts.
 * The 3-second Promise.race timeout still caps the total wall-clock time,
 * so hanging connections are detected promptly regardless of retries.
 */
export async function isDatabaseHealthy(): Promise<boolean> {
  if (!pool) {
    return false;
  }

  let timeoutHandle: NodeJS.Timeout | undefined;
  try {
    await Promise.race([
      pool.query("SELECT 1"),
      new Promise((_, reject) => {
        timeoutHandle = setTimeout(
          () => reject(new Error("Health check timeout")),
          3000,
        );
      }),
    ]);
    return true;
  } catch (error) {
    logger.warn({ error }, "Database health check failed");
    return false;
  } finally {
    if (timeoutHandle) {
      clearTimeout(timeoutHandle);
    }
  }
}

/**
 * Default export to use a Proxy to defer access until after database initialization.
 */
export default new Proxy({} as ReturnType<typeof getDb>, {
  get(_, prop) {
    return (getDb() as unknown as Record<string | symbol, unknown>)[prop];
  },
});

export { schema };

/**
 * Set the database instance directly (for testing purposes only).
 * This bypasses the normal initialization flow.
 *
 * Pass `null` on test-file teardown: between files in a shared worker,
 * accesses through the proxy then fail with getDb()'s "Database not
 * initialized" — a condition import-time consumers already tolerate —
 * instead of "PGlite is closed" from a torn-down instance.
 * @public — consumed via dynamic import in src/test/setup.ts
 */
export function __setTestDb(
  testDb: ReturnType<typeof drizzle<typeof schema>> | null,
): void {
  db = testDb;
}

// ============================================================
// Internal implementation
// ============================================================

let pool: pg.Pool | null = null;
let db: ReturnType<typeof drizzle<typeof schema>> | null = null;
let initPromise: Promise<void> | null = null;

async function doInitializeDatabase(): Promise<void> {
  let connectionString: string;

  const vaultRef = process.env[DATABASE_URL_VAULT_REF_ENV];
  if (vaultRef && isReadonlyVaultEnabled()) {
    // READONLY_VAULT is enabled and vault ref is set - read from Vault
    const vaultUrl = await getDatabaseUrlFromVault(vaultRef);
    if (vaultUrl) {
      logger.info(
        { connectionStringPrefix: vaultUrl.slice(0, 10) },
        "Database URL successfully loaded from Vault",
      );
      connectionString = vaultUrl;
    } else {
      logger.info("Database URL not found in Vault, falling back to env var");
      connectionString = config.database.url;
    }
  } else {
    // Use env var
    logger.info(
      "ARCHESTRA_DATABASE_URL_VAULT_REF is not set or READONLY_VAULT is not enabled, falling back to env var",
    );
    connectionString = config.database.url;
  }

  // Assign the globals only once everything has succeeded: a throw from the
  // instrumentation below must not leave a half-initialized db that makes
  // the next initializeDatabase() call return early instead of retrying.
  const newPool = createPool(connectionString);
  try {
    const newDb = drizzle({
      client: newPool,
      schema,
    });

    instrumentDrizzleClient(newDb, { dbSystem: "postgresql" });
    installDbErrorSafetyNet();

    pool = newPool;
    db = newDb;
  } catch (error) {
    await newPool.end().catch(() => {});
    throw error;
  }
  logger.info(
    { poolMax: config.database.poolMax },
    "Database connection pool initialized",
  );
}

/**
 * Create a connection pool with proper keepalive settings to prevent
 * "Connection terminated unexpectedly" errors.
 *
 * This addresses an issue where connections were being
 * terminated by network infrastructure (load balancers, NAT gateways) due to
 * idle timeouts.
 *
 * Pool configuration:
 * - max: ARCHESTRA_DATABASE_POOL_MAX (default 50). Must be coordinated with
 *   Postgres `max_connections` so that pods × poolMax stays under the server limit.
 * - idleTimeoutMillis: 30s (close idle connections after 30s)
 * - connectionTimeoutMillis: 10s (fail if can't get connection in 10s)
 * - statement_timeout: ARCHESTRA_DATABASE_STATEMENT_TIMEOUT_MILLIS (default 30s);
 *   kills runaway queries so a single slow query can't hang a connection forever.
 *
 * Connection keepalive configuration:
 * - keepAlive: true (enable TCP keepalive probes)
 * - keepAliveInitialDelayMillis: 10s (start probes after 10s of idle)
 *
 * The keepalive settings help prevent load balancers and NAT gateways from
 * terminating idle connections, which is a common cause of the
 * "Connection terminated unexpectedly" error in cloud environments.
 */
function createPool(connectionString: string): pg.Pool {
  const newPool = new pg.Pool({
    connectionString,
    // Pool configuration
    max: config.database.poolMax,
    idleTimeoutMillis: 30000,
    connectionTimeoutMillis: 10000,
    // Per-connection statement timeout: kills runaway queries instead of
    // letting them hang a connection indefinitely. 0 disables it.
    statement_timeout: config.database.statementTimeoutMillis,
    // Keepalive configuration to prevent "Connection terminated unexpectedly"
    keepAlive: true,
    keepAliveInitialDelayMillis: 10000,
  });

  /**
   * Handle errors on idle clients in the pool.
   * Without this handler, connection errors on idle clients would cause
   * an unhandled 'error' event and crash the process.
   * The pool will automatically remove the errored client and create a new one.
   */
  newPool.on("error", (err) => {
    logger.error({ err }, "Unexpected error on idle database client");
  });

  /**
   * Wrap pool.query() with automatic retry for transient connection errors
   * (ECONNREFUSED, connection timeout, connection terminated, etc.).
   *
   * This is safe because pool.query() checks out a fresh connection for each
   * call, so retries will get a new connection from the pool.
   * Transaction queries (via a checked-out PoolClient) are NOT affected.
   */
  wrapPoolWithRetry(newPool);

  return newPool;
}
