import { type SelectQueryBuilder, type ReferenceExpression, type OperandValueExpression, sql } from 'kysely'

type WhereInput<DB, TB extends keyof DB> = Partial<{
  [C in keyof DB[TB]]:
    | OperandValueExpression<DB, TB, DB[TB][C]>
    | readonly OperandValueExpression<DB, TB, DB[TB][C]>[]
    | null
}>

export function applyWhere<DB, TB extends keyof DB, O>(
  builder: SelectQueryBuilder<DB, TB, O>,
  where: WhereInput<DB, TB>,
): SelectQueryBuilder<DB, TB, O> {
  return Object.entries(where).reduce((builder, [key, value]) => {
    const column = key as ReferenceExpression<DB, TB>

    if (value === null) {
      return builder.where(column, 'is', null)
    } else if (Array.isArray(value)) {
      return builder.where(column, 'in', value)
    }

    return builder.where(column, '=', value)
  }, builder)
}

interface PaginationInput<DB, TB extends keyof DB> {
  where: WhereInput<DB, TB>
  limit?: number
  offset?: number
  orderBy?: {
    column: ReferenceExpression<DB, TB>
    direction?: 'asc' | 'desc'
  }
}

export function applyPagination<DB, TB extends keyof DB, O>(
  builder: SelectQueryBuilder<DB, TB, O>,
  { limit, offset, orderBy }: PaginationInput<DB, TB>,
): SelectQueryBuilder<DB, TB, O> {
  let qb = builder

  if (orderBy) {
    qb = qb.orderBy(orderBy.column, orderBy.direction ?? 'asc')
  }

  if (limit !== undefined) {
    qb = qb.limit(limit)
  }

  if (offset !== undefined) {
    qb = qb.offset(offset)
  }

  return qb
}

export async function paginate<DB, TB extends keyof DB, O>(
  baseDataQuery: SelectQueryBuilder<DB, TB, O>,
  baseCountQuery: SelectQueryBuilder<DB, TB, unknown>,
  input: PaginationInput<DB, TB>,
) {
  const dataQuery = applyPagination(input.where ? applyWhere(baseDataQuery, input.where) : baseDataQuery, input)

  let countQuery = baseCountQuery.select(sql<number>`count(*)::int`.as('totalCount'))

  if (input.where) {
    countQuery = applyWhere(countQuery, input.where)
  }

  const [countRow, items] = await Promise.all([countQuery.executeTakeFirstOrThrow(), dataQuery.execute()])

  return {
    items,
    totalCount: (countRow as { totalCount: number }).totalCount,
    limit: input.limit,
    offset: input.offset,
    count: items.length,
  }
}