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 { getProjectDataDir } from '../lib/project-data-dir'
import { getRootPath } from '../lib/get-root-path'
import type { TrpcContext } from '../lib/trpc-context'
import { analyzeTraceQuotes, generateSuggestedUserResponse } from '../tests/analyze'
import type { AgentState } from './agents'
import { fsContextPersistence, fsStatePersistence } from './fs-persistence'
import { generateReplyVariants, replaceSavedState } from './variants'

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
					}
				})
			}),
		generateUserResponse: t.procedure
			.input(
				z.object({
					model: z.string().trim().min(1).max(2048).optional(),
					stringifiedConversation: z.string(),
					response: z.string()
				})
			)
			.mutation(async ({ input }) => {
				return generateSuggestedUserResponse({
					model: input.model,
					stringifiedConversation: input.stringifiedConversation,
					response: input.response
				})
			}),
		savePrompt: t.procedure
			.input(
				z.object({
					agentId: AgentIdSchema,
					step: z.string().trim().min(1).max(512),
					prompt: z.string().max(1_000_000)
				})
			)
			.mutation(async ({ input }) => {
				await persistPromptUpdate(input.agentId, input.step, input.prompt)
				return { ok: true }
			}),
		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)
			}),
		generateVariants: t.procedure
			.input(
				z.object({
					agentId: AgentIdSchema,
					region: z.string().trim().min(1).max(128).nullable().optional(),
					prevStateId: z.string().trim().min(1).max(2048).nullable(),
					rawUserMessage: z.string().trim().min(1).max(100_000),
					context: z.record(z.string(), z.unknown()).nullable().optional(),
					promptUpdate: z
						.object({
							step: z.string().trim().min(1).max(512),
							prompt: z.string().max(1_000_000)
						})
						.nullable()
						.optional(),
					count: z.number().int().min(1).max(9).default(9)
				})
			)
			.mutation(async ({ input }) => {
				if (input.promptUpdate) {
					await persistPromptUpdate(
						input.agentId,
						input.promptUpdate.step,
						input.promptUpdate.prompt
					)
				}
				const stateKey = resolveAgentStateKey(input.agentId, input.region ?? null)
				const rawStates = await fsStatePersistence.load(stateKey)
				const states = Array.isArray(rawStates) ? (rawStates as AgentState[]) : []
				const context =
					input.context && typeof input.context === 'object'
						? input.context
						: ((await fsContextPersistence.load('default')) as Record<string, any>) || {}

				const variants = await generateReplyVariants({
					agentId: input.agentId,
					region: input.region ?? null,
					states,
					prevStateId: input.prevStateId ?? null,
					rawUserMessage: input.rawUserMessage,
					count: input.count,
					context
				})

				return variants.map((variant, index) => ({
					id: `variant-${index + 1}`,
					text: variant.text,
					state: {
						history: variant.state.history,
						schedule: variant.state.schedule
					}
				}))
			}),
		commitVariant: t.procedure
			.input(
				z.object({
					agentId: AgentIdSchema,
					region: z.string().trim().min(1).max(128).nullable().optional(),
					targetStateId: z.string().trim().min(1).max(2048),
					state: z.object({
						history: z.unknown(),
						schedule: z.unknown()
					})
				})
			)
			.mutation(async ({ input }) => {
				const stateKey = resolveAgentStateKey(input.agentId, input.region ?? null)
				const rawStates = await fsStatePersistence.load(stateKey)
				const states = Array.isArray(rawStates) ? (rawStates as AgentState[]) : []
				if (!states.length) throw new Error('No saved roleplay states found for this agent')
				if (states[states.length - 1]?.id !== input.targetStateId) {
					throw new Error('Variants can only be applied to the latest saved reply')
				}

				const nextStates = replaceSavedState(states, input.targetStateId, {
					history: input.state.history as AgentState['history'],
					schedule: input.state.schedule as AgentState['schedule']
				})
				await fsStatePersistence.save(stateKey, nextStates)
				return { ok: true }
			}),
		context: t.procedure.input(z.object({ agentId: z.string() })).query(async ({ input }) => {
			const savedContext = await fsContextPersistence.load('default')
			if (!savedContext) return {}
			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 resolveAgentStateKey(agentId: string, region: string | null): string {
	return region ? `${agentId}.${region}` : agentId
}

function resolveAgentHistoryFile(agentId: string): string {
	const baseDir = getProjectDataDir('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 persistPromptUpdate(agentId: string, stepName: string, prompt: string) {
	const { factory, basePath } = await getAgentFactory(agentId)
	const ctx = createContext({})
	const steps: any[] = []
	const stepRegistry = {
		addStep: (step: any) => {
			steps.push(step)
		},
		getSteps: () => steps
	}

	await factory({
		ctx,
		getMessages: () => [],
		logger: console,
		scheduler: {
			clearAllPendingActions: noopAsync,
			registerAction: () => noopAsync
		},
		sendAction: noopAsync,
		sendMessage: noopAsync,
		stepRegistry,
		tracer: stepRegistry,
		stepTracer: {
			addStepTrace: noop,
			flush: noopAsync
		}
	})

	const targetStep = steps.find((step) => step?.name === stepName)
	if (!targetStep || targetStep?.prompt?.type !== 'file') return

	const promptPath = path.resolve(basePath, targetStep.prompt.path)
	await fs.promises.writeFile(promptPath, prompt, 'utf-8')
}

async function getAgentFactory(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}`)
	}
	return { factory, basePath: agentFactories.basePath }
}

async function getPromptMapper(agentId: string) {
	const { factory } = await getAgentFactory(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()
}
