import type { TRPCRootObject } from '@trpc/server/unstable-core-do-not-import'
import { z } from 'zod'

import {
	addTest,
	deleteTest,
	getResultEntry,
	loadConfig,
	loadLatestResultsByTest,
	loadTests,
	patchTest,
	patchTestsBatch,
	reorderTests,
	RunConfigSchema,
	saveConfig
} from './data'
import { readSuite, TestRowSchema } from './suite'
import { cancelPromptTests, runPromptTestsManaged } from './academy'
import { analyzeRunQuotes } from './analyze'
import { findPromptIssues } from './prompt-issues'
import type { TrpcContext } from '../lib/trpc-context'

export function createTestsRouter(t: TRPCRootObject<TrpcContext, any, any>) {
	return t.router({
		list: t.procedure.input(z.object({ agent: z.string() })).query(({ input }) => {
			return loadTests(input.agent)
		}),
		config: t.procedure.input(z.object({ agent: z.string() })).query(({ input }) => {
			return loadConfig(input.agent)
		}),
		updateConfig: t.procedure
			.input(z.object({ agent: z.string(), cfg: RunConfigSchema }))
			.mutation(({ input }) => {
				return saveConfig(input.agent, input.cfg)
			}),
		latestResult: t.procedure.input(z.object({ agent: z.string() })).query(({ input }) => {
			return loadLatestResultsByTest(input.agent)
		}),
		resultEntry: t.procedure
			.input(
				z.object({
					agent: z.string(),
					testId: z.string(),
					// TODO resultKey should be keyof of ResultEntry
					resultKey: z.string(),
					index: z.number().optional(),
					timestamp: z.string().optional()
				})
			)
			.query(({ input }) => {
				return getResultEntry(
					input.agent,
					input.testId,
					input.resultKey as any,
					input.index,
					input.timestamp
				)
			}),
		analyzeRun: t.procedure
			.input(
				z.object({
					agent: z.string(),
					testId: z.string(),
					index: z.number().int().nonnegative(),
					timestamp: z.string().optional()
				})
			)
			.mutation(async ({ input }) => {
				return analyzeRunQuotes(input)
			}),
		findPromptIssues: t.procedure
			.input(
				z.object({
					agent: z.string(),
					step: z.string(),
					prompt: z.string()
				})
			)
			.mutation(async ({ input }) => {
				return findPromptIssues(input)
			}),
		suite: t.procedure.input(z.object({ agent: z.string() })).query(({ input }) => {
			return readSuite(input.agent)
		}),
		addTest: t.procedure
			.input(
				z.object({ agent: z.string(), mergeContexts: z.boolean(), seed: TestRowSchema.partial() })
			)
			.mutation(({ input }) => {
				return addTest(input.agent, input.mergeContexts, input.seed)
			}),
		patchTest: t.procedure
			.input(z.object({ agent: z.string(), id: z.string(), patch: TestRowSchema.partial() }))
			.mutation(({ input }) => {
				return patchTest(input.agent, input.id, input.patch)
			}),
		patchTestsBatch: t.procedure
			.input(
				z.object({
					agent: z.string(),
					updates: z.array(z.object({ id: z.string(), patch: TestRowSchema.partial() }))
				})
			)
			.mutation(({ input }) => {
				return patchTestsBatch(input.agent, input.updates)
			}),
		deleteTest: t.procedure
			.input(z.object({ agent: z.string(), id: z.string() }))
			.mutation(({ input }) => {
				return deleteTest(input.agent, input.id)
			}),
		reorder: t.procedure
			.input(z.object({ agent: z.string(), order: z.array(z.string()) }))
			.mutation(({ input }) => {
				return reorderTests(input.agent, input.order)
			}),
		runTests: t.procedure
			.input(
				z.union([
					z
						.object({
							agent: z.string(),
							scope: z.enum(['pinned', 'all']).optional()
						})
						.strict(),
					z
						.object({
							agent: z.string(),
							testId: z.string()
						})
						.strict()
				])
			)
			.mutation(({ input }) => {
				if ('testId' in input) {
					return runPromptTestsManaged(input.agent, { testId: input.testId })
				}
				return runPromptTestsManaged(input.agent, { scope: input.scope ?? 'pinned' })
			}),
		cancelRun: t.procedure.input(z.object({ agent: z.string() })).mutation(({ input }) => {
			return cancelPromptTests(input.agent)
		})
	})
}
