import type { RuntimeLogger } from "../../../../core/types.ts";
import type { MigrationDb } from "../runtime-migration-runner.ts";
import type { DialectAdapter, IdempotencyInsertValues, RuntimeDialect } from "./dialect-adapter.ts";
import type { NodePgDatabase } from "drizzle-orm/node-postgres";
import type pg from "pg";

import { sql } from "drizzle-orm";
import { drizzle } from "drizzle-orm/node-postgres";
import { Pool } from "pg";
import { runRuntimeMigrations } from "../runtime-migration-runner.ts";
import * as schema from "../schema/postgres.ts";
import { resolveDrizzleSession, runWithDrizzleSession } from "./session.ts";

export interface PostgresDialectOptions {
  connectionString: string;
  poolMax?: number;
  logger?: RuntimeLogger;
}

export interface PostgresDialect extends RuntimeDialect {
  readonly kind: "postgres";
  readonly db: NodePgDatabase<typeof schema>;
  readonly schema: typeof schema;
  readonly pool: pg.Pool;
}

const migrationDirectory = new URL("../../../../../drizzle/postgres/", import.meta.url);

export async function createPostgresDialect(options: PostgresDialectOptions): Promise<PostgresDialect> {
  const pool = new Pool({
    connectionString: options.connectionString,
    max: options.poolMax ?? 10,
  });

  try {
    // Hold one client for the entire migration run so BEGIN/COMMIT share a connection.
    const migrationClient = await pool.connect();
    try {
      await runRuntimeMigrations({
        db: wrapPgMigrationClient(migrationClient),
        migrationDirectory,
        logger: options.logger,
        dialect: "postgres",
      });
    } finally {
      migrationClient.release();
    }
  } catch (error) {
    await pool.end();
    throw error;
  }

  const rootDb = drizzle(pool, { schema });
  const dialect: PostgresDialect = {
    kind: "postgres",
    get db(): NodePgDatabase<typeof schema> {
      return resolveDrizzleSession(rootDb);
    },
    schema,
    adapter: null as unknown as DialectAdapter,
    pool,
    async close(): Promise<void> {
      await pool.end();
    },
  };
  dialect.adapter = createPostgresAdapter(rootDb, () => dialect.db);

  return dialect;
}

function createPostgresAdapter(
  rootDb: NodePgDatabase<typeof schema>,
  getDb: () => NodePgDatabase<typeof schema>,
): DialectAdapter {
  return {
    kind: "postgres",

    async transaction<T>(work: () => Promise<T>): Promise<T> {
      return await rootDb.transaction(async (tx) => {
        return await runWithDrizzleSession(tx, work);
      });
    },

    async deleteRetentionTail(keepCount: number): Promise<void> {
      await getDb().execute(sql`
        delete from runs
        where id in (
          select id from runs
          order by started_at desc, id desc
          offset ${keepCount}
        )
      `);
    },

    async insertIgnoreIdempotency(values: IdempotencyInsertValues): Promise<boolean> {
      const rows = await getDb()
        .insert(schema.idempotencyRecords)
        .values({
          keyHash: values.keyHash,
          claimId: values.claimId,
          requestHash: values.requestHash,
          state: "in_progress",
          responseValue: null,
          createdAt: values.createdAt,
          expiresAt: values.expiresAt,
        })
        .onConflictDoNothing()
        .returning({ keyHash: schema.idempotencyRecords.keyHash });
      return rows.length > 0;
    },
  };
}

function wrapPgMigrationClient(client: pg.PoolClient): MigrationDb {
  return {
    async exec(sqlText: string): Promise<void> {
      await client.query(sqlText);
    },
    prepare(sqlText: string) {
      const pgSql = toPgPlaceholders(sqlText);
      return {
        async run(...params: unknown[]): Promise<{ changes: number }> {
          const result = await client.query(pgSql, params);
          return { changes: result.rowCount ?? 0 };
        },
        async all(...params: unknown[]): Promise<unknown[]> {
          const result = await client.query(pgSql, params);
          return result.rows;
        },
        async get(...params: unknown[]): Promise<unknown> {
          const result = await client.query(pgSql, params);
          return result.rows[0];
        },
      };
    },
  };
}

/** Convert `?` placeholders used by the shared migration runner into `$1`, `$2`, … */
function toPgPlaceholders(sqlText: string): string {
  let index = 0;
  return sqlText.replace(/\?/g, () => {
    index += 1;
    return `$${index}`;
  });
}