import fs from 'fs'
import path from 'path'
import { z } from 'zod'
import { createContext } from '@recombine-ai/engine'
import type { TRPCRootObject } from '@trpc/server/unstable-core-do-not-import'
import { getAgentFactories } from '../lib/get-agent-factories'
import { getRootPath } from '../lib/get-root-path'
import type { TrpcContext } from '../lib/trpc-context'
import { analyzeTraceQuotes } from '../tests/analyze'
import { fsContextPersistence } from './fs-persistence'

const AgentIdSchema = z.string().trim().min(1).max(512)
const promptMappersCache = new Map<string, (ctx: object) => object>()

export function createRolePlayRouter(t: TRPCRootObject<TrpcContext, any, any>) {
	return t.router({
		list: t.procedure.query(async () => {
			const rootPath = await getRootPath()
			const agentFactories = await getAgentFactories(rootPath)
			return Object.keys(agentFactories.factories)
		}),
		analyzeTrace: t.procedure
			.input(
				z.object({
					agentId: AgentIdSchema,
					traceId: z.string().trim().min(1).max(2048).optional(),
					eventName: z.string().trim().min(1).max(256).optional(),
					createdAt: z.number().finite().optional(),
					model: z.string().trim().min(1).max(512).optional(),
					renderedPrompt: z.string(),
					stringifiedConversation: z.string(),
					response: z.string(),
					stepTrace: z.unknown().optional()
				})
			)
			.mutation(async ({ input }) => {
				return analyzeTraceQuotes({
					model: input.model,
					renderedPrompt: input.renderedPrompt,
					stringifiedConversation: input.stringifiedConversation,
					response: input.response,
					stepTrace: input.stepTrace as any,
					rawInput: {
						source: 'roleplay',
						agentId: input.agentId,
						traceId: input.traceId ?? null,
						eventName: input.eventName ?? null,
						createdAt: input.createdAt ?? null
					}
				})
			}),
		userMessageHistory: t.procedure
			.input(z.object({ agentId: AgentIdSchema }))
			.query(async ({ input }) => {
				const filePath = resolveAgentHistoryFile(input.agentId)
				const states = await loadJsonArray(filePath)

				const byText = new Map<string, { text: string; lastSeen: number; uses: number }>()
				let idx = 0

				for (const state of states) {
					if (!state || typeof state !== 'object') continue
					const history = (state as any).history
					if (!Array.isArray(history)) continue

					for (const entry of history) {
						if (!entry || typeof entry !== 'object') continue
						if ((entry as any).type !== 'message') continue
						const payload = (entry as any).payload
						if (!payload || typeof payload !== 'object') continue
						if ((payload as any).sender !== 'user') continue

						const rawText = (payload as any).text
						if (typeof rawText !== 'string') continue
						const text = rawText.trim()
						if (!text) continue

						idx += 1
						const existing = byText.get(text)
						if (existing) {
							existing.lastSeen = idx
							existing.uses += 1
						} else {
							byText.set(text, { text, lastSeen: idx, uses: 1 })
						}
					}
				}

				const items = Array.from(byText.values()).sort((a, b) => b.lastSeen - a.lastSeen)
				return items.slice(0, 2000)
			}),
		context: t.procedure.input(z.object({ agentId: z.string() })).query(async ({ input }) => {
			const savedContext = await fsContextPersistence.load('default')
			let mapper = promptMappersCache.get(input.agentId)
			if (mapper) {
				return mapper(savedContext)
			}

			mapper = await getPromptMapper(input.agentId)
			promptMappersCache.set(input.agentId, mapper)
			return mapper(savedContext)
		})
	})
}

function resolveAgentHistoryFile(agentId: string): string {
	const baseDir = path.resolve('.', 'data', 'history')
	const filePath = path.resolve(baseDir, `${agentId}.json`)

	// Prevent path traversal and absolute paths.
	if (filePath === baseDir || !filePath.startsWith(baseDir + path.sep)) {
		throw new Error('Invalid agentId')
	}

	return filePath
}

async function loadJsonArray(filePath: string): Promise<unknown[]> {
	try {
		const raw = await fs.promises.readFile(filePath, 'utf-8')
		const trimmed = raw.trim()
		if (!trimmed) return []
		const parsed = JSON.parse(trimmed) as unknown
		return Array.isArray(parsed) ? parsed : []
	} catch (e) {
		if ((e as any)?.code === 'ENOENT') return []
		console.error('[roleplay] Failed to load history file:', filePath, e)
		return []
	}
}

const noop = () => void 0
const noopAsync = () => Promise.resolve()

async function getPromptMapper(agentId: string) {
	const rootPath = await getRootPath()
	const agentFactories = await getAgentFactories(rootPath)
	const factory = agentFactories.factories[agentId]
	if (!factory) {
		throw new Error(`Cannot find agent ${agentId}`)
	}
	const ctx = createContext({})
	// create noop agent just to get promptContextMapper out of it... ಠ_ಠ
	await factory({
		ctx,
		getMessages: () => [],
		logger: console,
		scheduler: {
			clearAllPendingActions: noopAsync,
			registerAction: () => noopAsync
		},
		sendAction: noopAsync,
		sendMessage: noopAsync,
		stepRegistry: {
			addStep: () => {}
		},
		tracer: {
			addStep: () => {}
		},
		stepTracer: {
			addStepTrace: noop,
			flush: noopAsync
		}
	})

	return ctx.getPromptContextMapper()
}
