import type { TRPCRootObject } from '@trpc/server/unstable-core-do-not-import'
import { z } from 'zod'
import {
	addTest,
	deleteTest,
	getResultEntry,
	loadConfig,
	loadLatestResultsByTest,
	loadTests,
	patchTest,
	reorderTests,
	RunConfigSchema,
	saveConfig
} from './data'
import { readSuite, TestRowSchema } from './suite'
import { runPromptTests } from './academy'
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
				)
			}),
		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)
			}),
		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.object({ agent: z.string(), scope: z.enum(['pinned', 'all']).optional() }))
			.mutation(({ input }) => {
				return runPromptTests(input.agent, input.scope)
			})
	})
}
