/* eslint-disable sort-keys-fix/sort-keys-fix  */
import { TRPCError } from '@trpc/server';
import dayjs from 'dayjs';
import JSONL from 'jsonl-parse-stringify';
import pMap from 'p-map';
import { z } from 'zod';

import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
import { serverDB } from '@/database/server';
import { FileModel } from '@/database/server/models/file';
import {
  EvalDatasetModel,
  EvalDatasetRecordModel,
  EvalEvaluationModel,
  EvaluationRecordModel,
} from '@/database/server/models/ragEval';
import { authedProcedure, router } from '@/libs/trpc';
import { keyVaults } from '@/libs/trpc/middleware/keyVaults';
import { S3 } from '@/server/modules/S3';
import { createAsyncServerClient } from '@/server/routers/async';
import { getFullFileUrl } from '@/server/utils/files';
import {
  EvalDatasetRecord,
  EvalEvaluationStatus,
  InsertEvalDatasetRecord,
  RAGEvalDataSetItem,
  insertEvalDatasetRecordSchema,
  insertEvalDatasetsSchema,
  insertEvalEvaluationSchema,
} from '@/types/eval';

const ragEvalProcedure = authedProcedure.use(keyVaults).use(async (opts) => {
  const { ctx } = opts;

  return opts.next({
    ctx: {
      datasetModel: new EvalDatasetModel(ctx.userId),
      fileModel: new FileModel(serverDB, ctx.userId),
      datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
      evaluationModel: new EvalEvaluationModel(ctx.userId),
      evaluationRecordModel: new EvaluationRecordModel(ctx.userId),
      s3: new S3(),
    },
  });
});

export const ragEvalRouter = router({
  createDataset: ragEvalProcedure
    .input(
      z.object({
        description: z.string().optional(),
        knowledgeBaseId: z.string(),
        name: z.string(),
      }),
    )
    .mutation(async ({ input, ctx }) => {
      const data = await ctx.datasetModel.create({
        description: input.description,
        knowledgeBaseId: input.knowledgeBaseId,
        name: input.name,
      });

      return data?.id;
    }),

  getDatasets: ragEvalProcedure
    .input(z.object({ knowledgeBaseId: z.string() }))

    .query(async ({ ctx, input }): Promise<RAGEvalDataSetItem[]> => {
      return ctx.datasetModel.query(input.knowledgeBaseId);
    }),

  removeDataset: ragEvalProcedure
    .input(z.object({ id: z.number() }))
    .mutation(async ({ input, ctx }) => {
      return ctx.datasetModel.delete(input.id);
    }),

  updateDataset: ragEvalProcedure
    .input(
      z.object({
        id: z.number(),
        value: insertEvalDatasetsSchema.partial(),
      }),
    )
    .mutation(async ({ input, ctx }) => {
      return ctx.datasetModel.update(input.id, input.value);
    }),

  // Dataset Item operations
  createDatasetRecords: ragEvalProcedure
    .input(
      z.object({
        datasetId: z.number(),
        question: z.string(),
        ideal: z.string().optional(),
        referenceFiles: z.array(z.string()).optional(),
        metadata: z.record(z.unknown()).optional(),
      }),
    )
    .mutation(async ({ input, ctx }) => {
      const data = await ctx.datasetRecordModel.create(input);
      return data?.id;
    }),

  getDatasetRecords: ragEvalProcedure
    .input(z.object({ datasetId: z.number() }))
    .query(async ({ ctx, input }): Promise<EvalDatasetRecord[]> => {
      return ctx.datasetRecordModel.query(input.datasetId);
    }),

  removeDatasetRecords: ragEvalProcedure
    .input(z.object({ id: z.number() }))
    .mutation(async ({ input, ctx }) => {
      return ctx.datasetRecordModel.delete(input.id);
    }),

  updateDatasetRecords: ragEvalProcedure
    .input(
      z.object({
        id: z.number(),
        value: z
          .object({
            question: z.string(),
            ideal: z.string(),
            referenceFiles: z.array(z.string()),
            metadata: z.record(z.unknown()),
          })
          .partial(),
      }),
    )
    .mutation(async ({ input, ctx }) => {
      return ctx.datasetRecordModel.update(input.id, input.value);
    }),

  importDatasetRecords: ragEvalProcedure
    .input(
      z.object({
        datasetId: z.number(),
        pathname: z.string(),
      }),
    )
    .mutation(async ({ input, ctx }) => {
      const dataStr = await ctx.s3.getFileContent(input.pathname);
      const items = JSONL.parse<InsertEvalDatasetRecord>(dataStr);

      insertEvalDatasetRecordSchema.array().parse(items);

      const data = await Promise.all(
        items.map(async ({ referenceFiles, question, ideal }) => {
          const files = typeof referenceFiles === 'string' ? [referenceFiles] : referenceFiles;

          let fileIds: string[] | undefined = undefined;

          if (files) {
            const items = await ctx.fileModel.findByNames(files);

            fileIds = items.map((item) => item.id);
          }

          return {
            question,
            ideal,
            referenceFiles: fileIds,
            datasetId: input.datasetId,
          };
        }),
      );

      return ctx.datasetRecordModel.batchCreate(data);
    }),

  // Evaluation operations
  startEvaluationTask: ragEvalProcedure
    .input(z.object({ id: z.number() }))
    .mutation(async ({ input, ctx }) => {
      // Start evaluation task
      const evaluation = await ctx.evaluationModel.findById(input.id);

      if (!evaluation) {
        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
      }

      // create evaluation records by dataset records
      const datasetRecords = await ctx.datasetRecordModel.findByDatasetId(evaluation.datasetId);

      if (datasetRecords.length === 0) {
        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Dataset record is empty' });
      }

      const evalRecords = await ctx.evaluationRecordModel.batchCreate(
        datasetRecords.map((record) => ({
          evaluationId: input.id,
          datasetRecordId: record.id,
          question: record.question!,
          ideal: record.ideal,
          status: EvalEvaluationStatus.Pending,
          embeddingModel: DEFAULT_EMBEDDING_MODEL,
          languageModel: DEFAULT_MODEL,
        })),
      );

      const asyncCaller = await createAsyncServerClient(ctx.userId, ctx.jwtPayload);

      await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Processing });
      try {
        await pMap(
          evalRecords,
          async (record) => {
            asyncCaller.ragEval.runRecordEvaluation
              .mutate({ evalRecordId: record.id })
              .catch(async (e) => {
                await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Error });

                throw new TRPCError({
                  code: 'BAD_GATEWAY',
                  message: `[ASYNC_TASK] Failed to start evaluation task: ${e.message}`,
                });
              });
          },
          {
            concurrency: 30,
          },
        );

        return { success: true };
      } catch (e) {
        console.error('[startEvaluationTask]:', e);

        await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Error });

        return { success: false };
      }
    }),

  checkEvaluationStatus: ragEvalProcedure
    .input(z.object({ id: z.number() }))
    .query(async ({ input, ctx }) => {
      const evaluation = await ctx.evaluationModel.findById(input.id);

      if (!evaluation) {
        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
      }

      const records = await ctx.evaluationRecordModel.findByEvaluationId(input.id);

      const isSuccess = records.every((record) => record.status === EvalEvaluationStatus.Success);

      if (isSuccess) {
        // 将结果上传到 S3

        const evalRecords = records.map((record) => ({
          question: record.question,
          context: record.context,
          answer: record.answer,
          ground_truth: record.ideal,
        }));
        const date = dayjs().format('YYYY-MM-DD-HH-mm');
        const filename = `${date}-eval_${evaluation.id}-${evaluation.name}.jsonl`;
        const path = `rag_eval_records/${filename}`;

        await ctx.s3.uploadContent(path, JSONL.stringify(evalRecords));

        // 保存数据
        await ctx.evaluationModel.update(input.id, {
          status: EvalEvaluationStatus.Success,
          evalRecordsUrl: await getFullFileUrl(path),
        });
      }

      return { success: isSuccess };
    }),
  createEvaluation: ragEvalProcedure
    .input(insertEvalEvaluationSchema)
    .mutation(async ({ input, ctx }) => {
      const data = await ctx.evaluationModel.create({
        description: input.description,
        knowledgeBaseId: input.knowledgeBaseId,
        datasetId: input.datasetId,
        name: input.name,
      });

      return data?.id;
    }),

  removeEvaluation: ragEvalProcedure
    .input(z.object({ id: z.number() }))
    .mutation(async ({ input, ctx }) => {
      return ctx.evaluationModel.delete(input.id);
    }),

  getEvaluationList: ragEvalProcedure
    .input(z.object({ knowledgeBaseId: z.string() }))
    .query(async ({ ctx, input }) => {
      return ctx.evaluationModel.queryByKnowledgeBaseId(input.knowledgeBaseId);
    }),
});