/* eslint-disable @typescript-eslint/no-loop-func */

import { setImmediate } from "node:timers/promises";
import { TRANSACTION_ROLLBACK_ERROR_PREFIX } from "../constants.js";
import { transactionContext } from "../contexts/transactionContext.js";
import type { ConnectionPoolClient } from "../factories/createConnectionPool.js";
import { getPoolClientState } from "../state.js";
import type {
  ClientConfiguration,
  Logger,
  Query,
  QueryContext,
  QueryId,
  QueryResult,
  QueryResultRow,
  StreamResult,
} from "../types.js";
import { getStackTrace } from "./getStackTrace.js";
import { SpanStatusCode, trace } from "@opentelemetry/api";
import type { DriverNotice } from "@slonik/driver";
import {
  BackendTerminatedError,
  BackendTerminatedUnexpectedlyError,
  DataIntegrityError,
  InvalidInputError,
  NotFoundError,
  SlonikError,
  TupleMovedToAnotherPartitionError,
  UnexpectedForeignConnectionError,
  UnexpectedStateError,
} from "@slonik/errors";
import type { PrimitiveValueExpression, QuerySqlToken } from "@slonik/sql-tag";
import { generateUid } from "@slonik/utilities";
import PQueue from "p-queue";
import { serializeError } from "serialize-error";

export type IntegrityValidation = {
  validationType:
    | "MANY_COLUMNS"
    | "MANY_ROWS"
    | "MANY_ROWS_ONE_COLUMN"
    | "MAYBE_MANY_ROWS_KEY_VALUE_COLUMNS"
    | "MAYBE_MANY_ROWS_ONE_COLUMN"
    | "MAYBE_ONE_COLUMN"
    | "MAYBE_ONE_ROW"
    | "ONE_COLUMN"
    | "ONE_ROW";
};

const tracer = trace.getTracer("slonik.interceptors");

export type ExecutionRoutine = (
  connection: ConnectionPoolClient,
  sql: string,
  values: readonly PrimitiveValueExpression[],
  queryContext: QueryContext,
  query: Query,
) => Promise<GenericQueryResult>;

type GenericQueryResult = QueryResult<QueryResultRow> | StreamResult;

type TransactionQuery = {
  readonly executionContext: QueryContext;
  readonly executionRoutine: ExecutionRoutine;
  readonly name?: string;
  readonly sql: string;
  readonly values: readonly PrimitiveValueExpression[];
};

const retryQuery = async (
  connectionLogger: Logger,
  connection: ConnectionPoolClient,
  query: TransactionQuery,
  retryLimit: number,
) => {
  let result: GenericQueryResult;
  let remainingRetries = retryLimit;
  let attempt = 0;

  // @todo Provide information about the queries being retried to the logger.
  while (remainingRetries-- > 0) {
    attempt++;

    try {
      connectionLogger.trace(
        {
          attempt,
          queryId: query.executionContext.queryId,
        },
        "retrying query",
      );

      result = await query.executionRoutine(
        connection,
        query.sql,

        // @todo Refresh execution context to reflect that the query has been re-tried.
        query.values,

        // This (probably) requires changing `queryId` and `queryInputTime`.
        // It should be needed only for the last query (because other queries will not be processed by the middlewares).
        query.executionContext,
        {
          name: query.name,
          sql: query.sql,
          values: query.values,
        },
      );

      // If the attempt succeeded break out of the loop
      break;
    } catch (error) {
      if (
        typeof error.code === "string" &&
        error.code.startsWith(TRANSACTION_ROLLBACK_ERROR_PREFIX) &&
        remainingRetries > 0
      ) {
        continue;
      }

      throw error;
    }
  }

  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  return result!;
};

type StackCrumb = {
  columnNumber: null | number;
  fileName: null | string;
  functionName: null | string;
  lineNumber: null | number;
};

// oxlint-disable-next-line complexity
const executeQueryInternal = async (
  connectionLogger: Logger,
  connection: ConnectionPoolClient,
  clientConfiguration: ClientConfiguration,
  query: QuerySqlToken,
  inheritedQueryId: QueryId | undefined,
  executionRoutine: ExecutionRoutine,
  integrityValidation?: IntegrityValidation,
  stream: boolean = false,
): Promise<
  QueryResult<Record<string, PrimitiveValueExpression>> | StreamResult
  // eslint-disable-next-line complexity
> => {
  const poolClientState = getPoolClientState(connection);

  if (poolClientState.terminated) {
    throw new BackendTerminatedError(poolClientState.terminated);
  }

  // Convert Slonik's internal `$slonik_N` placeholders to `$N` once, up front, so the
  // input-validation guards below see real Postgres placeholders. In particular, a query
  // that is only a value binding arrives here as `$slonik_1` and must normalize to `$1`
  // for the guard to fire. See the comment in `formatSlonikPlaceholder` for why the
  // internal prefix exists.
  const normalizedSql = query.sql.replaceAll("$slonik_", "$");
  const trimmedSql = normalizedSql.trim();

  if (trimmedSql === "") {
    throw new InvalidInputError("Unexpected SQL input. Query cannot be empty.");
  }

  if (trimmedSql === "$1") {
    throw new InvalidInputError(
      "Unexpected SQL input. Query cannot be empty. Found only value binding.",
    );
  }

  const transactionStore = transactionContext.getStore();

  if (
    clientConfiguration.dangerouslyAllowForeignConnections !== true &&
    transactionStore?.transactionId &&
    transactionStore.transactionId !== poolClientState.transactionId
  ) {
    throw new UnexpectedForeignConnectionError();
  }

  const queryInputTime = process.hrtime.bigint();

  let stackTrace: null | StackCrumb[] = null;

  if (clientConfiguration.captureStackTrace) {
    stackTrace = getStackTrace();
  }

  const queryId = inheritedQueryId ?? generateUid();

  const log = connectionLogger.child({
    queryId,
  });

  const originalQuery = {
    // Include statement name for prepared statements if provided
    name: query.name,
    sql: normalizedSql,
    values: query.values,
  };

  let actualQuery: Query = {
    name: originalQuery.name,
    sql: originalQuery.sql,
    values: originalQuery.values,
  };

  const executionContext: QueryContext = {
    connectionId: poolClientState.connectionId,
    log,
    originalQuery,
    poolId: poolClientState.poolId,
    queryId,
    queryInputTime,
    resultParser: query.parser,
    sandbox: {},
    stackTrace,
    transactionId: poolClientState.transactionId,
  };

  for (const interceptor of clientConfiguration.interceptors) {
    const beforeTransformQuery = interceptor.beforeTransformQuery;

    if (beforeTransformQuery) {
      if (clientConfiguration.tracing) {
        await tracer.startActiveSpan("slonik.interceptor.beforeTransformQuery", async (span) => {
          span.setAttribute("interceptor.name", interceptor.name);

          try {
            await beforeTransformQuery(executionContext, actualQuery);
          } catch (error) {
            span.recordException(error);
            span.setStatus({
              code: SpanStatusCode.ERROR,
              message: String(error),
            });

            throw error;
          } finally {
            span.end();
          }
        });
      } else {
        await beforeTransformQuery(executionContext, actualQuery);
      }
    }
  }

  for (const interceptor of clientConfiguration.interceptors) {
    const transformQuery = interceptor.transformQuery;

    if (transformQuery) {
      if (clientConfiguration.tracing) {
        actualQuery = tracer.startActiveSpan("slonik.interceptor.transformQuery", (span) => {
          span.setAttribute("interceptor.name", interceptor.name);

          try {
            return transformQuery(executionContext, actualQuery);
          } catch (error) {
            span.recordException(error);
            span.setStatus({
              code: SpanStatusCode.ERROR,
              message: String(error),
            });

            throw error;
          } finally {
            span.end();
          }
        });
      } else {
        actualQuery = transformQuery(executionContext, actualQuery);
      }
    }
  }

  let result: GenericQueryResult | null;

  if (!stream) {
    for (const interceptor of clientConfiguration.interceptors) {
      const beforeQueryExecution = interceptor.beforeQueryExecution;

      if (beforeQueryExecution) {
        if (clientConfiguration.tracing) {
          result = await tracer.startActiveSpan(
            "slonik.interceptor.beforeQueryExecution",
            {
              attributes: {
                "interceptor.name": interceptor.name,
              },
            },
            async (span) => {
              try {
                return await beforeQueryExecution(executionContext, actualQuery);
              } catch (error) {
                span.recordException(error);
                span.setStatus({
                  code: SpanStatusCode.ERROR,
                  message: String(error),
                });

                throw error;
              } finally {
                span.end();
              }
            },
          );
        } else {
          result = await beforeQueryExecution(executionContext, actualQuery);
        }

        if (result) {
          log.info(
            "beforeQueryExecution interceptor produced a result; short-circuiting query execution using beforeQueryExecution result",
          );

          return result;
        }
      }
    }
  }

  const notices: DriverNotice[] = [];

  const noticeListener = (notice: DriverNotice) => {
    notices.push(notice);
  };

  const activeQuery = Promise.withResolvers<null>();

  const blockingPromise = poolClientState.activeQuery?.promise ?? null;

  poolClientState.activeQuery = activeQuery;

  await blockingPromise;

  connection.on("notice", noticeListener);

  const queryWithContext = {
    executionContext,
    executionRoutine,
    name: actualQuery.name,
    sql: actualQuery.sql,
    values: actualQuery.values,
  };

  try {
    try {
      try {
        // eslint-disable-next-line require-atomic-updates
        result = await executionRoutine(
          connection,
          actualQuery.sql,
          actualQuery.values,
          executionContext,
          actualQuery,
        );
      } catch (error) {
        const shouldRetry =
          typeof error.code === "string" &&
          error.code.startsWith(TRANSACTION_ROLLBACK_ERROR_PREFIX) &&
          clientConfiguration.queryRetryLimit > 0;

        // Transactions errors in queries that are part of a transaction are handled by the transaction/nestedTransaction functions
        if (shouldRetry && !poolClientState.transactionId) {
          // eslint-disable-next-line require-atomic-updates
          result = await retryQuery(
            connectionLogger,
            connection,
            queryWithContext,
            clientConfiguration.queryRetryLimit,
          );
        } else {
          throw error;
        }
      }
    } catch (error) {
      // The driver is responsible for throwing an appropriately wrapped error.
      if (error instanceof BackendTerminatedError) {
        poolClientState.terminated = error;
      }

      // If the error has been already handled by the driver, then we should not wrap it again.
      if (!(error instanceof SlonikError)) {
        if (error.message === "Connection terminated unexpectedly") {
          throw new BackendTerminatedUnexpectedlyError(error);
        }

        if (
          error.message.includes(
            "tuple to be locked was already moved to another partition due to concurrent update",
          )
        ) {
          throw new TupleMovedToAnotherPartitionError(error);
        }
      }

      error.notices = notices;

      throw error;
    } finally {
      connection.off("notice", noticeListener);

      activeQuery.resolve(null);
    }
  } catch (error) {
    log.error(
      {
        error: serializeError(error),
      },
      "execution routine produced an error",
    );

    for (const interceptor of clientConfiguration.interceptors) {
      if (interceptor.queryExecutionError) {
        await interceptor.queryExecutionError(executionContext, actualQuery, error, notices);
      }
    }

    error.notices = notices;

    throw error;
  }

  if (!result) {
    throw new UnexpectedStateError("Expected query result to be returned.");
  }

  // @ts-expect-error -- We want to keep notices as readonly for consumer, but write to it here.
  result.notices = notices;

  // `clientConfiguration.interceptors` is `readonly` and never mutated during a
  // query (the before/error hooks above already iterate it directly), so we read
  // it as-is instead of copying it on every query.
  const interceptors = clientConfiguration.interceptors;

  if (result.type !== "QueryResult") {
    return result;
  }

  try {
    if (integrityValidation) {
      if (integrityValidation.validationType === "ONE_ROW") {
        if (result.rows.length === 0) {
          throw new NotFoundError("Query returned no rows.", actualQuery);
        }

        if (result.rows.length > 1) {
          throw new DataIntegrityError("Query returned multiple rows.", actualQuery);
        }
      }

      if (integrityValidation.validationType === "MAYBE_ONE_ROW" && result.rows.length > 1) {
        throw new DataIntegrityError("Query returned multiple rows.", actualQuery);
      }

      // Cache the column count from the first row to avoid repeated Object.keys() calls.
      // It is safe to assume that whatever the first row is, it will be the same for all rows.
      const firstRowColumnCount = result.rows.length > 0 ? Object.keys(result.rows[0]).length : 0;

      if (integrityValidation.validationType === "ONE_COLUMN") {
        if (result.rows.length === 0) {
          throw new NotFoundError("Query returned no rows.", actualQuery);
        }

        if (result.rows.length !== 1) {
          throw new DataIntegrityError("Query returned multiple rows.", actualQuery);
        }

        if (firstRowColumnCount !== 1) {
          throw new DataIntegrityError("Query returned rows with multiple columns.", actualQuery);
        }
      }

      if (integrityValidation.validationType === "MAYBE_ONE_COLUMN") {
        if (result.rows.length > 1) {
          throw new DataIntegrityError("Query returned multiple rows.", actualQuery);
        }

        if (result.rows.length === 1 && firstRowColumnCount !== 1) {
          throw new DataIntegrityError("Query returned rows with multiple columns.", actualQuery);
        }
      }

      if (integrityValidation.validationType === "MANY_ROWS" && result.rows.length === 0) {
        throw new NotFoundError("Query returned no rows.", actualQuery);
      }

      if (integrityValidation.validationType === "MANY_ROWS_ONE_COLUMN") {
        if (result.rows.length === 0) {
          throw new NotFoundError("Query returned no rows.", actualQuery);
        }

        if (firstRowColumnCount !== 1) {
          throw new DataIntegrityError("Query returned rows with multiple columns.", actualQuery);
        }
      }

      if (
        integrityValidation.validationType === "MAYBE_MANY_ROWS_ONE_COLUMN" &&
        result.rows.length > 0 &&
        firstRowColumnCount !== 1
      ) {
        throw new DataIntegrityError("Query returned rows with multiple columns.", actualQuery);
      }

      if (
        integrityValidation.validationType === "MAYBE_MANY_ROWS_KEY_VALUE_COLUMNS" &&
        result.rows.length > 0
      ) {
        const firstRowColumnNames = Object.keys(result.rows[0]);

        if (
          firstRowColumnNames.length !== 2 ||
          !firstRowColumnNames.includes("key") ||
          !firstRowColumnNames.includes("value")
        ) {
          throw new DataIntegrityError(
            'Query returned rows with columns other than "key" and "value".',
            actualQuery,
          );
        }

        const seenKeys = new Set<string>();

        for (const row of result.rows) {
          // Keys are stringified because that is how they will be represented in the resulting record.
          const key = String((row as Record<string, unknown>).key);

          if (seenKeys.has(key)) {
            throw new DataIntegrityError("Query returned duplicate keys.", actualQuery);
          }

          seenKeys.add(key);
        }
      }
    }
  } catch (error) {
    for (const interceptor of clientConfiguration.interceptors) {
      if (interceptor.dataIntegrityError) {
        await interceptor.dataIntegrityError(executionContext, actualQuery, error, result);
      }
    }

    throw error;
  }

  for (const interceptor of interceptors) {
    const afterQueryExecution = interceptor.afterQueryExecution;

    if (afterQueryExecution) {
      if (clientConfiguration.tracing) {
        await tracer.startActiveSpan(
          "slonik.interceptor.afterQueryExecution",
          {
            attributes: {
              "interceptor.name": interceptor.name,
            },
          },
          async (span) => {
            try {
              await afterQueryExecution(
                executionContext,
                actualQuery,
                result as QueryResult<QueryResultRow>,
              );
            } catch (error) {
              span.recordException(error);
              span.setStatus({
                code: SpanStatusCode.ERROR,
                message: String(error),
              });

              throw error;
            } finally {
              span.end();
            }
          },
        );
      } else {
        await afterQueryExecution(
          executionContext,
          actualQuery,
          result as QueryResult<QueryResultRow>,
        );
      }
    }
  }

  for (const interceptor of interceptors) {
    const transformRow = interceptor.transformRow;

    if (transformRow) {
      const { fields, rows } = result;

      let transformedRows: QueryResultRow[];

      if (clientConfiguration.tracing) {
        transformedRows = tracer.startActiveSpan(
          "slonik.interceptor.transformRow",
          {
            attributes: {
              "interceptor.name": interceptor.name,
              "rows.length": rows.length,
            },
          },
          (span) => {
            try {
              return rows.map((row) => {
                return transformRow(executionContext, actualQuery, row, fields);
              });
            } catch (error) {
              span.recordException(error);
              span.setStatus({
                code: SpanStatusCode.ERROR,
                message: String(error),
              });

              throw error;
            } finally {
              span.end();
            }
          },
        );
      } else {
        transformedRows = rows.map((row) => {
          return transformRow(executionContext, actualQuery, row, fields);
        });
      }

      // avoid spreading the result object to avoid performance overhead
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (result as any).rows = transformedRows;
    }
  }

  for (const interceptor of interceptors) {
    const transformRowAsync = interceptor.transformRowAsync;

    if (transformRowAsync) {
      const { fields, rows } = result;

      let transformedRows: QueryResultRow[];

      const processRowsAsync = async (): Promise<QueryResultRow[]> => {
        const queue = new PQueue({ concurrency: 10 });
        const results: QueryResultRow[] = [];
        let firstError: unknown;

        for (let i = 0; i < rows.length; i++) {
          void queue
            .add(async () => {
              // Yield to the macrotask queue periodically so parallel queries
              // can interleave their row processing instead of monopolizing the event loop.
              if (i > 0 && i % 100 === 0) {
                await setImmediate();
              }

              results[i] = await transformRowAsync(executionContext, actualQuery, rows[i], fields);
            })
            .catch((error) => {
              if (!firstError) {
                firstError = error;
                queue.clear();
              }
            });
        }

        await queue.onIdle();

        if (firstError) {
          throw firstError;
        }

        return results;
      };

      if (clientConfiguration.tracing) {
        transformedRows = await tracer.startActiveSpan(
          "slonik.interceptor.transformRowAsync",
          {
            attributes: {
              "interceptor.name": interceptor.name,
              "rows.length": rows.length,
            },
          },
          async (span) => {
            try {
              return await processRowsAsync();
            } catch (error) {
              span.recordException(error);
              span.setStatus({
                code: SpanStatusCode.ERROR,
                message: String(error),
              });

              throw error;
            } finally {
              span.end();
            }
          },
        );
      } else {
        transformedRows = await processRowsAsync();
      }

      // avoid spreading the result object to avoid performance overhead
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (result as any).rows = transformedRows;
    }
  }

  for (const interceptor of interceptors) {
    const beforeQueryResult = interceptor.beforeQueryResult;

    if (beforeQueryResult) {
      if (clientConfiguration.tracing) {
        await tracer.startActiveSpan(
          "slonik.interceptor.beforeQueryResult",
          {
            attributes: {
              "interceptor.name": interceptor.name,
            },
          },
          async (span) => {
            try {
              await beforeQueryResult(
                executionContext,
                actualQuery,
                result as QueryResult<QueryResultRow>,
              );
            } catch (error) {
              span.recordException(error);
              span.setStatus({
                code: SpanStatusCode.ERROR,
                message: String(error),
              });

              throw error;
            } finally {
              span.end();
            }
          },
        );
      } else {
        await beforeQueryResult(
          executionContext,
          actualQuery,
          result as QueryResult<QueryResultRow>,
        );
      }
    }
  }

  return result;
};

export const executeQuery = async (
  connectionLogger: Logger,
  connection: ConnectionPoolClient,
  clientConfiguration: ClientConfiguration,
  query: QuerySqlToken,
  inheritedQueryId: QueryId | undefined,
  executionRoutine: ExecutionRoutine,
  stream: boolean,
  integrityValidation?: IntegrityValidation,
): Promise<QueryResult<Record<string, PrimitiveValueExpression>> | StreamResult> => {
  if (!clientConfiguration.tracing) {
    return await executeQueryInternal(
      connectionLogger,
      connection,
      clientConfiguration,
      query,
      inheritedQueryId,
      executionRoutine,
      integrityValidation,
      stream,
    );
  }

  return await tracer.startActiveSpan(
    "slonik.executeQuery",
    {
      attributes: {
        sql: query.sql,
      },
    },
    async (span) => {
      try {
        return await executeQueryInternal(
          connectionLogger,
          connection,
          clientConfiguration,
          query,
          inheritedQueryId,
          executionRoutine,
          integrityValidation,
          stream,
        );
      } catch (error) {
        span.recordException(error);
        span.setStatus({
          code: SpanStatusCode.ERROR,
          message: String(error),
        });

        throw error;
      } finally {
        span.end();
      }
    },
  );
};