import type { PgDatabase, PgQueryResultHKT, PgTransaction, PgTransactionConfig } from 'drizzle-orm/pg-core'
import type { ExtractTablesWithRelations, SQLWrapper } from 'drizzle-orm'
import type { TransactionCapability } from './transaction.js'

export {
  assertTransactionIdentity,
  getTransactionIdentity,
  isTransactionCapabilityError,
  TransactionCapabilityError,
  withTransactionIdentity,
} from './transaction.js'
export type { TransactionIdentity } from './transaction.js'

/** The host's drizzle schema (pgTables + relations). db core defines none — it is generic over this. */
export type Schema = Record<string, unknown>

/** Query-capable handle WITHOUT interactive transactions. What `./postgres/neon-http` returns. */
export type Database<S extends Schema = Record<string, never>> =
  Omit<PgDatabase<PgQueryResultHKT, S, ExtractTablesWithRelations<S>>, 'transaction'>

/** Dialect-neutral callback-only mutation handle. */
export type Transaction<S extends Schema = Record<string, never>> =
  TransactionCapability<S> & {
    execute<Row extends Record<string, unknown> = Record<string, unknown>>(
      query: SQLWrapper,
    ): Promise<readonly Row[]>
  }

type PgTransactionHandle<S extends Schema> = Pick<
  PgTransaction<PgQueryResultHKT, S, ExtractTablesWithRelations<S>>,
  | 'query'
  | '$with'
  | '$count'
  | 'with'
  | 'select'
  | 'selectDistinct'
  | 'selectDistinctOn'
  | 'update'
  | 'insert'
  | 'delete'
  | 'refreshMaterializedView'
>

/** Strongly typed PostgreSQL handle supplied by PostgreSQL adapter callbacks. */
export type PostgresTransaction<S extends Schema = Record<string, never>> =
  PgTransactionHandle<S> &
  Transaction<S> & {
    transaction<T>(fn: (tx: PostgresTransaction<S>) => Promise<T>, config?: PgTransactionConfig): Promise<T>
  }

/** Adds interactive transactions back. What `./postgres/neon-serverless`, `./postgres/postgres-js`, `./postgres/pglite` return. */
export type TransactionalDatabase<S extends Schema = Record<string, never>> = Database<S> & {
  transaction<T>(fn: (tx: PostgresTransaction<S>) => Promise<T>, config?: PgTransactionConfig): Promise<T>
}

/** The workhorse PostgreSQL query handle, in or out of a transaction. */
export type Querier<S extends Schema = Record<string, never>> = Database<S> | PostgresTransaction<S>

/** Dialect identifier for adapter selection */
export type Dialect = 'postgres' | 'sqlite'