import { dataWrapper, populatePk } from 'src/helpers/dbHelpers';
import {
  isAttachment,
  isLinksOrLTAR,
  NcApiVersion,
  type NcRequest,
} from 'nocodb-sdk';
import { AttachmentUrlUploadPreparator } from './attachment-url-upload-preparator';
import type { Column } from 'src/models';
import type { IBaseModelSqlV2 } from '../IBaseModelSqlV2';
import { handleUniqueConstraintError } from '~/helpers/uniqueConstraintErrorHandler';

export const baseModelInsert = (baseModel: IBaseModelSqlV2) => {
  const single = async (
    data,
    request: NcRequest,
    trx?,
    _disableOptimization = false,
  ) => {
    let insertObj: Record<string, any> | undefined;
    try {
      const columns = await baseModel.model.getColumns(baseModel.context);
      const dbDataWrapper = dataWrapper(data);
      // exclude auto increment columns in body
      for (const col of columns) {
        if (col.ai) {
          const keyName = dbDataWrapper.getColumnKeyName(col);

          if (data[keyName]) {
            delete data[keyName];
          }
        }
      }

      await populatePk(baseModel.context, baseModel.model, data);

      // todo: filter based on view
      insertObj = await baseModel.model.mapAliasToColumn(
        baseModel.context,
        data,
        baseModel.clientMeta,
        baseModel.dbDriver,
        columns,
      );

      await baseModel.validate(insertObj, columns);

      if ('beforeInsert' in baseModel) {
        await baseModel.beforeInsert(insertObj, trx, request);
      }

      await baseModel.prepareNocoData(insertObj, true, request);

      let response;
      // const driver = trx ? trx : baseModel.dbDriver;

      const query = baseModel.dbDriver(baseModel.tnPath).insert(insertObj);
      if (baseModel.isPg && baseModel.model.primaryKey) {
        query.returning(
          `${baseModel.model.primaryKey.column_name} as ${baseModel.model.primaryKey.id}`,
        );
        response = await baseModel.execAndParse(query, null, { raw: true });
      }

      const ai = baseModel.model.columns.find((c) => c.ai);

      let ag: Column;
      if (!ai) ag = baseModel.model.columns.find((c) => c.meta?.ag);

      // handle if autogenerated primary key is used
      if (ag) {
        if (!response) {
          await baseModel.execAndParse(query);
        }
        response = await baseModel.readByPk(
          baseModel.extractCompositePK({
            rowId: insertObj[ag.column_name],
            insertObj,
            ag,
            ai,
          }),
          false,
          {},
          { ignoreView: true, getHiddenColumn: true },
        );
      } else if (
        !response ||
        (typeof response?.[0] !== 'object' && response?.[0] !== null)
      ) {
        let id;
        if (response?.length) {
          id = response[0];
        } else {
          const res = await baseModel.execAndParse(query, null, {
            raw: true,
          });
          id = res?.id ?? res[0]?.insertId ?? res;
        }

        if (ai) {
          if (baseModel.isSqlite) {
            // sqlite doesnt return id after insert
            id = (
              await baseModel.execAndParse(
                baseModel
                  .dbDriver(baseModel.tnPath)
                  .select(ai.column_name)
                  .max(ai.column_name, { as: '__nc_ai_id' }),
                null,
                { raw: true, first: true },
              )
            )?.__nc_ai_id;
          } else if (baseModel.isSnowflake || baseModel.isDatabricks) {
            id = (
              await baseModel.execAndParse(
                baseModel.dbDriver(baseModel.tnPath).max(ai.column_name, {
                  as: '__nc_ai_id',
                }),
                null,
                { raw: true, first: true },
              )
            ).__nc_ai_id;
          }
          response = await baseModel.readByPk(
            baseModel.extractCompositePK({ rowId: id, insertObj, ag, ai }),
            false,
            {},
            { ignoreView: true, getHiddenColumn: true },
          );
        } else {
          response = data;
        }
      } else if (ai) {
        const id = Array.isArray(response)
          ? response?.[0]?.[ai.id]
          : response?.[ai.id];
        response = await baseModel.readByPk(
          baseModel.extractCompositePK({ rowId: id, insertObj, ag, ai }),
          false,
          {},
          { ignoreView: true, getHiddenColumn: true },
        );
      }

      await baseModel.afterInsert({
        data: response,
        insertData: data,
        trx,
        req: request,
      });

      await baseModel.statsUpdate({
        count: 1,
      });

      return Array.isArray(response) ? response[0] : response;
    } catch (e: any) {
      // Handle unique constraint violations - this will throw if it's a unique constraint error
      await handleUniqueConstraintError({
        error: e,
        baseModel: baseModel,
        insertData: insertObj || data,
      });
      await baseModel.errorInsert(e, data, trx, request);
      throw e;
    }
  };
  const bulk = async (
    datas: any[],
    {
      chunkSize: _chunkSize = 100,
      cookie,
      foreign_key_checks = true,
      skip_hooks = false,
      raw = false,
      insertOneByOneAsFallback = false,
      isSingleRecordInsertion = false,
      typecast = false,
      allowSystemColumn = false,
      undo = false,
      apiVersion = NcApiVersion.V2,
    }: {
      chunkSize?: number;
      cookie?: NcRequest;
      foreign_key_checks?: boolean;
      skip_hooks?: boolean;
      raw?: boolean;
      insertOneByOneAsFallback?: boolean;
      isSingleRecordInsertion?: boolean;
      allowSystemColumn?: boolean;
      typecast?: boolean;
      undo?: boolean;
      apiVersion?: NcApiVersion;
    } = {},
  ) => {
    let trx;
    try {
      const insertDatas = raw ? datas : [];
      const postInsertOpsMap: Record<
        number,
        ((rowId: any) => Promise<string>)[]
      > = {};
      let preInsertOps: (() => Promise<string>)[] = [];
      let aiPkCol: Column;
      let agPkCol: Column;
      let columns: Column[];

      if (!raw) {
        columns = await baseModel.model.getColumns(baseModel.context);

        const order = await baseModel.getHighestOrderInTable();
        const nestedCols = columns.filter((c) => isLinksOrLTAR(c));
        const attachmentCols = columns.filter((c) => isAttachment(c));

        for (const [index, d] of datas.entries()) {
          const insertObj = await baseModel.handleValidateBulkInsert(
            d,
            columns,
            {
              allowSystemColumn,
              undo,
              typecast,
            },
          );

          await baseModel.prepareNocoData(insertObj, true, cookie, null, {
            ncOrder: order?.plus(index),
            undo,
          });

          // prepare nested link data for insert only if it is single record insertion
          if (isSingleRecordInsertion || apiVersion === NcApiVersion.V3) {
            const operations = await baseModel.prepareNestedLinkQb({
              nestedCols,
              data: d,
              insertObj,
              req: cookie,
            });

            postInsertOpsMap[index] = operations.postInsertOps;
            preInsertOps = operations.preInsertOps;
          }
          if (attachmentCols.length > 0) {
            const attachmentOperations =
              await new AttachmentUrlUploadPreparator().prepareAttachmentUrlUpload(
                baseModel,
                {
                  attachmentCols,
                  data: insertObj,
                  req: cookie,
                },
              );
            postInsertOpsMap[index] = [
              ...(postInsertOpsMap[index] ?? []),
              ...(attachmentOperations.postInsertOps ?? []),
            ];
            preInsertOps = [].concat(
              ...(preInsertOps ?? []),
              ...(attachmentOperations.preInsertOps ?? []),
            );
          }

          insertDatas.push(insertObj);
        }

        aiPkCol = baseModel.model.primaryKeys.find((pk) => pk.ai);
        agPkCol = baseModel.model.primaryKeys.find((pk) => pk.meta?.ag);
      } else {
        columns = await baseModel.model.getColumns(baseModel.context);

        const order = await baseModel.getHighestOrderInTable();

        await Promise.all(
          insertDatas.map(
            async (d, i) =>
              await baseModel.prepareNocoData(d, true, cookie, null, {
                raw,
                undo: undo,
                ncOrder: order?.plus(i),
              }),
          ),
        );
      }

      if ('beforeBulkInsert' in baseModel) {
        await baseModel.beforeBulkInsert(insertDatas, trx, cookie, {
          allowSystemColumn,
        });
      }

      // await baseModel.beforeInsertb(insertDatas, null);

      // fallbacks to `10` if database client is sqlite
      // to avoid `too many SQL variables` error
      // refer : https://www.sqlite.org/limits.html
      const chunkSize = baseModel.isSqlite ? 10 : _chunkSize;

      trx = await baseModel.dbDriver.transaction();

      if (!foreign_key_checks) {
        if (baseModel.isPg) {
          await trx.raw('set session_replication_role to replica;');
        } else if (baseModel.isMySQL) {
          await trx.raw('SET foreign_key_checks = 0;');
        }
      }

      await baseModel.runOps(
        preInsertOps.map((f) => f()),
        trx,
      );

      let responses;

      // insert one by one as fallback to get ids for sqlite and mysql
      if (
        insertOneByOneAsFallback &&
        (baseModel.isSqlite || baseModel.isMySQL)
      ) {
        // sqlite and mysql doesn't support returning, so insert one by one and return ids
        responses = [];

        for (const insertData of insertDatas) {
          const query = trx(baseModel.tnPath).insert(insertData);
          let id = (await query)[0];

          if (agPkCol) {
            id = insertData[agPkCol.column_name];
          }

          responses.push(
            baseModel.extractCompositePK({
              rowId: id,
              ai: aiPkCol,
              ag: agPkCol,
              insertObj: insertData,
              force: true,
            }) || insertData,
          );
        }
      } else {
        const returningObj: Record<string, string> = {};

        for (const col of baseModel.model.primaryKeys) {
          returningObj[col.title] = col.column_name;
        }

        responses =
          !raw && baseModel.isPg
            ? await trx
                .batchInsert(baseModel.tnPath, insertDatas, chunkSize)
                .returning(
                  baseModel.model.primaryKeys?.length ? returningObj : '*',
                )
            : await trx.batchInsert(baseModel.tnPath, insertDatas, chunkSize);
      }

      if (!foreign_key_checks) {
        if (baseModel.isPg) {
          await trx.raw('set session_replication_role to origin;');
        } else if (baseModel.isMySQL) {
          await trx.raw('SET foreign_key_checks = 1;');
        }
      }

      // insert nested link data for single record insertion or v3
      if (isSingleRecordInsertion || apiVersion === NcApiVersion.V3) {
        for (let i = 0; i < responses.length; i++) {
          const row = responses[i];
          let rowId = row[baseModel.model.primaryKey?.title];

          if (aiPkCol || agPkCol) {
            rowId = baseModel.extractCompositePK({
              rowId,
              ai: aiPkCol,
              ag: agPkCol,
              insertObj: insertDatas[i],
            });
          }

          await baseModel.runOps(
            postInsertOpsMap[i].map((f) => f(rowId)),
            trx,
          );
        }
      }

      await trx.commit();

      if (!raw && !skip_hooks) {
        if (isSingleRecordInsertion) {
          const insertData = await baseModel.readByPk(responses[0]);
          await baseModel.afterInsert({
            data: insertData,
            insertData: datas?.[0],
            trx: baseModel.dbDriver,
            req: cookie,
          });
        } else {
          await baseModel.afterBulkInsert(
            insertDatas.map((data, index) => {
              return {
                ...responses[index],
                ...data,
              };
            }),
            baseModel.dbDriver,
            cookie,
          );
        }
      }

      await baseModel.statsUpdate({
        count: insertDatas.length,
      });

      return responses;
    } catch (e: any) {
      await trx?.rollback();
      // Handle unique constraint violations - this will throw if it's a unique constraint error
      await handleUniqueConstraintError({
        error: e,
        baseModel: baseModel,
        insertData: datas?.[0],
      });
      // await baseModel.errorInsertb(e, data, null);
      throw e;
    }
  };
  return {
    single,
    bulk,
  };
};
