import 'dotenv/config'
import { WebSocketServer } from 'ws'
import { oldServer } from './old-server'
// ---- tRPC stuff
// ref: https://github.com/trpc/trpc/blob/main/examples/standalone-server/src/server.ts
import { initTRPC } from '@trpc/server'
import type { CreateHTTPContextOptions } from '@trpc/server/adapters/standalone'
import { createHTTPServer } from '@trpc/server/adapters/standalone'
import type { CreateWSSContextFnOptions } from '@trpc/server/adapters/ws'
import { z } from 'zod'
import { getRootPath } from './lib/get-root-path'
import { getAgentFactories } from './lib/get-agent-factories'
import {
	addTest,
	deleteTest,
	getResultEntry,
	loadConfig,
	loadLatestResultsByTest,
	loadTests,
	patchTest,
	reorderTests,
	RunConfigSchema,
	saveConfig
} from './lib/academy/data'
import { readSuite, TestRowSchema } from './lib/academy/suite'
import { runPromptTests } from './lib/academy/academy'

const WS_PORT = 8080
const HTTP_PORT = 2022

const wss = new WebSocketServer({ port: WS_PORT })

wss.on('connection', oldServer)

console.log(`WebSocket server started on ws://localhost:${WS_PORT}`)

// context is what shared between trpc calls or how you put external stuff in it (e.g. DB connections)
function createContext(opts: CreateHTTPContextOptions | CreateWSSContextFnOptions) {
	return {}
}
type Context = Awaited<ReturnType<typeof createContext>>

const t = initTRPC.context<Context>().create()

const publicProcedure = t.procedure
const router = t.router

const agentsRouter = router({
	list: publicProcedure.query(async () => {
		const rootPath = await getRootPath()
		const agentFactories = await getAgentFactories(rootPath)
		return Object.keys(agentFactories.factories)
	})
})

const testsRouter = router({
	list: publicProcedure.input(z.object({ agent: z.string() })).query(({ input }) => {
		return loadTests(input.agent)
	}),
	config: publicProcedure.input(z.object({ agent: z.string() })).query(({ input }) => {
		return loadConfig(input.agent)
	}),
	updateConfig: publicProcedure
		.input(z.object({ agent: z.string(), cfg: RunConfigSchema }))
		.mutation(({ input }) => {
			return saveConfig(input.agent, input.cfg)
		}),
	latestResult: publicProcedure.input(z.object({ agent: z.string() })).query(({ input }) => {
		return loadLatestResultsByTest(input.agent)
	}),
	resultEntry: publicProcedure
		.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: publicProcedure.input(z.object({ agent: z.string() })).query(({ input }) => {
		return readSuite(input.agent)
	}),
	addTest: publicProcedure
		.input(z.object({ agent: z.string(), seed: TestRowSchema.partial() }))
		.mutation(({ input }) => {
			return addTest(input.agent, input.seed)
		}),
	patchTest: publicProcedure
		.input(z.object({ agent: z.string(), id: z.string(), patch: TestRowSchema.partial() }))
		.mutation(({ input }) => {
			return patchTest(input.agent, input.id, input.patch)
		}),
	deleteTest: publicProcedure
		.input(z.object({ agent: z.string(), id: z.string() }))
		.mutation(({ input }) => {
			return deleteTest(input.agent, input.id)
		}),
	reorder: publicProcedure
		.input(z.object({ agent: z.string(), order: z.array(z.string()) }))
		.mutation(({ input }) => {
			return reorderTests(input.agent, input.order)
		}),
	runTests: publicProcedure
		.input(z.object({ agent: z.string(), scope: z.enum(['pinned', 'all']).optional() }))
		.mutation(({ input }) => {
			return runPromptTests(input.agent, input.scope)
		})
})

const appRouter = router({
	agents: agentsRouter,
	tests: testsRouter
})

export type AppRouter = typeof appRouter

const server = createHTTPServer({
	middleware(req, res, next) {
		res.setHeader('Access-Control-Allow-Origin', '*')
		res.setHeader('Access-Control-Request-Method', '*')
		res.setHeader('Access-Control-Allow-Methods', 'OPTIONS, GET, POST, PUT')
		res.setHeader('Access-Control-Allow-Headers', '*')
		if (req.method === 'OPTIONS') {
			res.writeHead(200)
			res.end()
			return
		}
		next()
	},
	router: appRouter,
	createContext
})

server.listen(HTTP_PORT)

console.log(`tRPC server started at http://localhost:${HTTP_PORT}`)

process.on('SIGINT', shutdown)
process.on('SIGTERM', shutdown)

function shutdown() {
	server.close()
	wss.close()
}
