import {
	createContext,
	makeAction,
	type Action,
	type ApiCallTrace,
	type ApiCallTracer,
	type ConversationalTrace,
	type ConversationalTracer,
	type Message,
	type StepTracer,
	type TestAgent,
	type TestVoiceAgent,
	type StepTrace
} from '@recombine-ai/engine'

import type { LogEntry } from '../../lib/ws-message'
import { getAgentFactories } from '../lib/get-agent-factories'
import { getImportNonce } from '../lib/import-nonce'
import { getRootPath } from '../lib/get-root-path'
import type { AgentState } from './agents'
import { createLogger } from './logger'
import { createScheduler } from './scheduler'
import { createStepRegistry } from './step-registry'

export type GeneratedVariant = {
	text: string
	state: AgentState
}

export async function generateReplyVariants(props: {
	agentId: string
	region?: string | null
	states: AgentState[]
	prevStateId: string | null
	rawUserMessage: string
	count: number
	context: Record<string, any>
}): Promise<GeneratedVariant[]> {
	const { agentId, region, states, prevStateId, rawUserMessage, count, context } = props
	if (!Number.isFinite(count) || count <= 0) return []

	const rootPath = await getRootPath()
	const { factories } = await getAgentFactories(rootPath, getImportNonce())
	const factory = factories[agentId]
	if (!factory) throw new Error(`Cannot find agent ${agentId}`)

	const statesById = new Map<string, AgentState>()
	for (const state of states) statesById.set(state.id, clone(state))

	const prevState = prevStateId ? (statesById.get(prevStateId) ?? null) : null
	if (prevStateId && !prevState) {
		throw new Error('Could not find the previous roleplay state for this reply')
	}

	const message = parseRawUserMessage(rawUserMessage)
	return Promise.all(
		Array.from({ length: count }, (_, index) =>
			generateSingleReplyVariant({
				index,
				factory,
				region,
				context,
				prevState,
				statesById,
				rawUserMessage,
				message
			})
		)
	)
}

export function replaceSavedState(
	states: AgentState[],
	targetStateId: string,
	replacement: Pick<AgentState, 'history' | 'schedule'>
): AgentState[] {
	const targetIndex = states.findIndex((state) => state.id === targetStateId)
	if (targetIndex === -1) throw new Error('Could not find the target state to replace')

	const target = states[targetIndex]
	const next = states.slice()
	next[targetIndex] = {
		...target,
		history: clone(replacement.history),
		schedule: clone(replacement.schedule)
	}
	return next
}

function parseRawUserMessage(raw: string): Message {
	const trimmed = String(raw ?? '').trim()
	if (!trimmed) throw new Error('Could not reconstruct the user message for variants')

	const imageMatch = /^\[image\]\s+(.+)$/u.exec(trimmed)
	if (imageMatch) {
		const imageUrl = (imageMatch[1] ?? '').trim()
		if (!imageUrl) throw new Error('Image variant message is missing its URL')
		return { sender: 'user', text: '', imageUrl }
	}

	return { sender: 'user', text: trimmed }
}

function collectMessagesUpToState(
	statesById: Map<string, AgentState>,
	stateId: string | null
): Message[] {
	if (!stateId) return []

	const chain: AgentState[] = []
	let currentId: string | null = stateId

	while (currentId) {
		const state = statesById.get(currentId)
		if (!state) break
		chain.push(state)
		currentId = state.prev
	}

	chain.reverse()

	const messages: Message[] = []
	for (const state of chain) {
		messages.push(...collectMessagesFromHistory(state.history))
	}
	return messages
}

function collectMessagesFromHistory(history: AgentState['history']): Message[] {
	const messages: Message[] = []
	for (const entry of history) {
		if (entry.type !== 'message') continue
		messages.push(entry.payload as Message)
	}
	return messages
}

function extractVariantText(state: AgentState): string {
	for (let index = state.history.length - 1; index >= 0; index--) {
		const entry = state.history[index]
		if (entry.type !== 'conversational-trace') continue
		const trace = entry.payload as ConversationalTrace
		if (trace.eventName === 'finished-main-streaming-step' || trace.eventName === 'sent-message') {
			return String(trace.content ?? '')
		}
	}

	for (let index = state.history.length - 1; index >= 0; index--) {
		const entry = state.history[index]
		if (entry.type !== 'message') continue
		const payload = entry.payload as Message
		if (payload.sender === 'agent') return String(payload.text ?? '')
	}

	return ''
}

function createVariantState(prev: string | null, index: number): AgentState {
	const now = new Date()
	return {
		id: `${now.toISOString()}::variant-${index + 1}-${Math.random().toString(36).slice(2, 8)}`,
		name: `${now.toLocaleString()} (Variant ${index + 1})`,
		history: [],
		schedule: {},
		prev,
		next: null
	}
}

async function generateSingleReplyVariant(props: {
	index: number
	factory: Awaited<ReturnType<typeof getAgentFactories>>['factories'][string]
	region: string | null | undefined
	context: Record<string, any>
	prevState: AgentState | null
	statesById: Map<string, AgentState>
	rawUserMessage: string
	message: Message
}): Promise<GeneratedVariant> {
	const { index, factory, region, context, prevState, statesById, rawUserMessage, message } = props

	const state = createVariantState(prevState?.id ?? null, index)
	const scheduleSeed = clone(prevState?.schedule ?? {})

	const log = (entry: LogEntry) => {
		state.history.push({
			type: 'log',
			payload: entry
		})
	}
	const logger = createLogger(log)

	const sendAction = async (action: Action) => {
		state.history.push({
			type: 'action',
			payload: action
		})
	}

	const scheduler = createScheduler({
		logger,
		sendAction,
		onUpdate: () => {}
	})
	const stepRegistry = createStepRegistry(logger)
	const stepTracer: StepTracer = {
		addStepTrace: (trace: StepTrace) => {
			state.history.push({
				type: 'trace',
				payload: trace
			})
		},
		async flush() {}
	}
	const apiCallTracer: ApiCallTracer = {
		addApiCallTrace: (trace: ApiCallTrace) => {
			state.history.push({
				type: 'api-call-trace',
				payload: trace
			})
		},
		async flush() {}
	}
	const conversationalTracer: ConversationalTracer = {
		addConversationalTrace: (trace: ConversationalTrace) => {
			;(trace as any).bosunStateId = state.id
			;(trace as any).bosunPrevStateId = state.prev
			;(trace as any).bosunUserMessage = rawUserMessage
			state.history.push({
				type: 'conversational-trace',
				payload: trace
			})
		},
		async flush() {}
	}

	const ctx = createContext({})
	if (region) {
		const current = (ctx.get() as Record<string, unknown> | null | undefined) || {}
		ctx.set({ ...current, recombineRegion: region })
	}
	if (context && typeof context === 'object') {
		ctx.swap(region ? { ...context, recombineRegion: region } : context)
	}

	const getMessages = () => [
		...collectMessagesUpToState(statesById, state.prev),
		...collectMessagesFromHistory(state.history)
	]

	const agent = await factory({
		stepTracer,
		tracer: stepRegistry,
		stepRegistry,
		logger,
		apiCallTracer,
		conversationalTracer,
		scheduler,
		getMessages,
		sendMessage: async (nextMessage: string | Message) => {
			const payload =
				typeof nextMessage === 'string'
					? ({ sender: 'agent', text: nextMessage } as Message)
					: nextMessage

			state.history.push({
				type: 'message',
				payload
			})
		},
		sendAction,
		ctx
	})

	scheduler.setState(scheduleSeed)

	if (!state.prev && hasStart(agent)) {
		await Promise.resolve().then(() => agent.start())
	}

	state.history.push({
		type: 'message',
		payload: message
	})

	if (isVoiceAgent(agent)) {
		await agent.streamResponse()
	} else if (await agent.isAssigned()) {
		try {
			const reactAct = makeAction(sendAction, 'Agent', 'React on message')
			await reactAct('started')
			await agent.reactOnMessage()
			await reactAct('completed')

			const respondAct = makeAction(sendAction, 'Agent', 'Compose reply')
			await respondAct('started')
			await agent.respondToMessage()
			await respondAct('completed')
		} catch (error) {
			if (typeof (agent as any)?.onFatalError === 'function') {
				await (agent as any).onFatalError(error)
			} else {
				throw error
			}
		}
	}

	state.schedule = clone(scheduler.serialize())
	return {
		text: extractVariantText(state),
		state
	}
}

function isVoiceAgent(agent: TestAgent): agent is TestVoiceAgent {
	return 'streamResponse' in agent
}

function hasStart(agent: TestAgent): agent is TestAgent & { start: () => Promise<void> | void } {
	return typeof (agent as { start?: unknown }).start === 'function'
}

function clone<T>(value: T): T {
	return JSON.parse(JSON.stringify(value))
}
