import {
	createContext,
	makeAction,
	type Action,
	type ApiCallTrace,
	type ApiCallTracer,
	type ConversationalTrace,
	type ConversationalTracer,
	type Message,
	type TestAgent,
	type TestAgentFactory,
	type TestVoiceAgent,
	type StepTracer,
	type StepTrace
} from '@recombine-ai/engine'
import type { DelayedActions, LogEntry } from '../../lib/ws-message'
import fs from 'fs'
import path, { join } from 'path'
import type { Context, ContextPersistence, StatePersistence } from './fs-persistence'
import { createLogger } from './logger'
import type { MessageDispatcher } from './message-dispatcher'
import { createScheduler, type BosunScheduler } from './scheduler'
import { createStepRegistry } from './step-registry'
import { getAgentFactories } from '../lib/get-agent-factories'
import { getRootPath } from '../lib/get-root-path'

interface MessageEntry {
	type: 'message'
	payload: Message
}
interface ActionEntry {
	type: 'action'
	payload: Action
}

interface HistoryLogEntry {
	type: 'log'
	payload: LogEntry
}

interface HistoryTraceEntry {
	type: 'trace'
	payload: StepTrace
}

interface HistoryApiCallTraceEntry {
	type: 'api-call-trace'
	payload: ApiCallTrace
}

interface HistoryConversationalTraceEntry {
	type: 'conversational-trace'
	payload: ConversationalTrace
}

type HistoryEntry =
	| MessageEntry
	| ActionEntry
	| HistoryLogEntry
	| HistoryTraceEntry
	| HistoryApiCallTraceEntry
	| HistoryConversationalTraceEntry

export interface AgentState {
	id: string
	name: string
	history: HistoryEntry[]
	schedule: DelayedActions
	prev: string | null
	next: string | null
}

type AgentWithState = Awaited<ReturnType<typeof loadAgent>>

const agentStates: Record<string, AgentState[]> = {}

export function createAgents(
	dispatcher: MessageDispatcher,
	statePersistence: StatePersistence,
	contextPersistence: ContextPersistence
) {
	let factories: Record<string, TestAgentFactory> = {}
	const agents: Record<string, AgentWithState> = {}
	let basePath = ''
	let currentAgent = ''

	async function getCurrentAgent() {
		if (!currentAgent) {
			currentAgent = Object.keys(factories)[0]
		}

		let agent = agents[currentAgent]
		if (!agent) {
			const factory = factories[currentAgent]
			if (!factory) {
				throw new Error(`Agent not found ${currentAgent}`)
			}

			const loadedStates = await statePersistence.load(currentAgent)
			agentStates[currentAgent] = loadedStates
			agent = await loadAgent({
				...dispatcher,
				agentId: currentAgent,
				contextPersistence,
				statePersistence,
				factory,
				basePath
			})
			agents[currentAgent] = agent
		}
		return agent
	}

	return {
		load: async () => {
			const rootPath = await getRootPath()
			const loaded = await getAgentFactories(rootPath)
			basePath = loaded.basePath
			factories = loaded.factories
			console.log(`createAgents, load; loaded agents: ${Object.keys(factories)}`)
			await dispatcher.sendLoaded(Object.keys(factories))
		},
		select: async (agentId: string, stateId?: string | null) => {
			console.log('createAgents, select; agents: ', Object.keys(agents), 'selected: ', agentId)
			currentAgent = agentId
			const agent = await getCurrentAgent()

			dispatcher.sendAgentSelected({
				id: agentId,
				steps: await Promise.all(
					agent.tracer.getSteps().map(async (s) => ({
						name: s.name,
						prompt: await s.prompt.content(),
						type: s.type,
						schema:
							s.prompt.type === 'file' && s.schema
								? await loadSchemaFile(agent.basePath, s.prompt.path)
								: undefined
					}))
				)
			})
			if (agent.ctx.get()) {
				const context = agent.ctx.get()
				const promptContext = agent.ctx.getPromptContextMapper()(context)
				dispatcher.sendContext(context, promptContext)
			}

			const states = agentStates[agentId] || []
			let isInitialLoad = false
			if (stateId === undefined && states.length > 0) {
				stateId = states[states.length - 1].id
				isInitialLoad = true
			}
			dispatcher.sendUpdateStates(
				states.map((state) => ({
					id: state.id,
					name: state.name,
					userMessage: userMessage(state)
				}))
			)
			await agent.selectState(states, stateId, isInitialLoad)
		},
		selectStep: async (step: string) => {
			await (await getCurrentAgent()).selectStep(step)
		},
		updateContext: async (context: Context) => {
			await (await getCurrentAgent()).updateContext(context)
		},
		updatePrompt: async (prompt: string) => {
			await (await getCurrentAgent()).updatePrompt(prompt)
		},
		triggerAction: async (action: string) => {
			await (await getCurrentAgent()).triggerAction(action)
		},
		handleMessage: async (message: Message) => {
			await dispatcher.sendMessage(message) // echo message
			await (await getCurrentAgent()).handleMessage(message)
		}
	}
}

type Public<T> = { [P in keyof T]: T[P] }
interface LoadProps extends Public<MessageDispatcher> {
	agentId: string
	basePath: string
	factory: TestAgentFactory
	contextPersistence: ContextPersistence
	statePersistence: StatePersistence
}

async function loadAgent(props: LoadProps) {
	let states: Record<string, AgentState> = {}
	let state = createState({}, null)
	let currentStep = ''
	let scheduler: BosunScheduler | null = null
	let agent: TestAgent | null = null
	const ctx = createContext({})
	ctx.subscribe((context) => {
		console.log('------ context updated --------')
		const promptContext = ctx.getPromptContextMapper()(context)
		props.sendContext(context, promptContext)
	})
	const logger = createLogger((msg) => {
		state.history.push({
			type: 'log',
			payload: msg
		})
		props.sendLog(msg)
	})

	const stepTracer: StepTracer = {
		addStepTrace: (stepTrace) => {
			state.history.push({
				type: 'trace',
				payload: stepTrace
			})
			props.sendTrace(stepTrace)
		},
		async flush() {}
	}

	const apiCallTracer: ApiCallTracer = {
		addApiCallTrace: (apiCallTrace) => {
			logger.log(
				'[apiCallTracer] addApiCallTrace',
				apiCallTrace?.operationName,
				apiCallTrace?.method,
				apiCallTrace?.path
			)
			state.history.push({
				type: 'api-call-trace',
				payload: apiCallTrace
			})
			props.sendApiCallTrace(apiCallTrace)
		},
		async flush() {}
	}

	const conversationalTracer: ConversationalTracer = {
		addConversationalTrace: (conversationalTrace) => {
			logger.log(
				'[conversationalTracer] addConversationalTrace',
				conversationalTrace?.eventName,
				conversationalTrace?.role
			)
			state.history.push({
				type: 'conversational-trace',
				payload: conversationalTrace
			})
			props.sendConversationalTrace(conversationalTrace)
		},
		async flush() {}
	}

	async function sendAction(action: Action) {
		state.history.push({
			type: 'action',
			payload: action
		})
		props.sendAction(action)
	}

	const tracer = createStepRegistry(logger)

	scheduler = createScheduler({
		logger,
		sendAction,
		onUpdate: async (actions) => {
			await props.sendSchedule(actions)
		}
	})

	agent = await props.factory({
		stepTracer: stepTracer,
		tracer,
		stepRegistry: tracer,
		logger,
		apiCallTracer,
		conversationalTracer,
		scheduler,
		getMessages: () => {
			let messages: Array<MessageEntry['payload']> = []
			let curState = state
			while (true) {
				curState.history
					.filter((entry) => entry.type === 'message')
					.reverse()
					.forEach((entry) => {
						messages.push(entry.payload)
					})
				if (!curState.prev) {
					break
				}
				curState = states[curState.prev]
			}
			return messages.reverse()
		},
		sendMessage: async (message: string | Message) => {
			if (typeof message === 'string') {
				state.history.push({
					type: 'message',
					payload: { text: message, sender: 'agent' }
				})
				await props.sendMessage({ text: message, sender: 'agent' })
			} else {
				state.history.push({
					type: 'message',
					payload: message
				})
				await props.sendMessage({ text: message.text, sender: message.sender })
			}
		},
		sendAction,
		ctx
	})

	// loading agent from factory loads it with default context.
	// if saved context is not empty, let's replace context in ctx with it.
	const savedContext = await props.contextPersistence.load('default')
	if (savedContext && Object.keys(savedContext).length > 0) {
		ctx.set(savedContext)
	} else if(ctx.get()) {
		await props.contextPersistence.save('default', ctx.get())
	}

	async function saveStates(states: AgentState[]) {
		await props.statePersistence.save(props.agentId, states)
		agentStates[props.agentId] = states
		console.log('------------- save states ---------------')
		states.forEach((state) => {
			console.log('State:', state.name, state.prev, state.next)
		})
		console.log('-----------------------------------------')

		await props.sendUpdateStates(
			states.map((state) => ({
				id: state.id,
				name: state.name,
				userMessage: userMessage(state)
			}))
		)
	}

	function addState() {
		if (!scheduler) {
			throw new Error('Cannot add state: agent is not initialized')
		}
		console.log('!! add state')
		const newState = createState(scheduler.serialize(), state.id)
		state.next = newState.id
		states[state.id] = clone({ ...state, schedule: scheduler.serialize() })
		saveStates(Object.values(states))
		state = newState
	}

	return {
		agent,
		tracer,
		basePath: props.basePath,
		ctx,
		selectState: async (
			stateHistory: AgentState[],
			stateId?: string | null,
			isInitialLoad?: boolean
		) => {
			states = {}
			stateHistory.forEach((state) => {
				states[state.id] = state
			})
			console.log('states: ' + stateHistory.length, stateId)
			if (stateHistory.length === 0 || !stateId) {
				// this agent loaded for the first time
				console.log(' --- init agent ---')
				state = createState({}, null)
				if (agent) {
					await (isVoiceAgent(agent) ? agent.onCallStart() : agent.start())
				}
				return
			}
			let loadedState = states[stateId]
			if (!loadedState) {
				return
			}
			console.log(' --- loaded state history: ', loadedState.history.length, '---')
			state = loadedState
			const replay: AgentState[] = []
			replay.push(loadedState)
			while (loadedState.prev) {
				loadedState = states[loadedState.prev]
				replay.push(loadedState)
			}
			console.log('replay: ' + replay.length)
			replay.reverse().forEach((state) => {
				state.history.forEach((entry) => {
					switch (entry.type) {
						case 'message':
							props.sendMessage(entry.payload)
							break
						case 'action':
							props.sendAction(entry.payload)
							break
						case 'log':
							props.sendLog(entry.payload)
							break
						case 'trace':
							props.sendTrace(entry.payload)
							break
						case 'api-call-trace':
							props.sendApiCallTrace(entry.payload)
							break
						case 'conversational-trace':
							props.sendConversationalTrace(entry.payload)
							break
					}
				})
			})
			scheduler?.setState(state.schedule)
			if (!isInitialLoad) {
				addState()
			}
		},
		selectStep: async (step: string) => {
			currentStep = step
			const selectedStep = tracer.getSteps().find((s) => s.name === step)
			const prompt = await selectedStep?.prompt.content()
			if (prompt) {
				await props.sendPrompt(prompt)
			}
		},
		updateContext: async (context: Context) => {
			ctx.swap(context)
			await props.contextPersistence.save('default', context)
		},
		updatePrompt: async (prompt: string) => {
			console.log('--- update prompt ---')
			const selectedStep = tracer.getSteps().find((s) => s.name === currentStep)
			if (selectedStep?.prompt.type !== 'file') {
				return
			}
			const promptPath = selectedStep?.prompt.path
			// TODO replace with persistence
			const path = join(props.basePath, selectedStep?.prompt.path)
			fs.promises
				.writeFile(path, prompt, 'utf-8')
				.then(() => {
					console.debug('FS: saved prompt:', promptPath)
				})
				.catch((error) => {
					console.error('FS: failed to save prompt:', error)
				})
		},
		triggerAction: async (action: string) => {
			if (!agent) {
				throw new Error('Cannot trigger action: agent is not initialized')
			}
			if (!scheduler) {
				throw new Error('Cannot trigger action: scheduler is not initialized')
			}
			if ('isAssigned' in agent && (await agent.isAssigned())) {
				try {
					await scheduler.triggerAction(action)
				} catch (e) {
					await agent.onFatalError(e as any)
				}
				state.schedule = scheduler.serialize()
			}
		},
		handleMessage: async (message: Message) => {
			if (!agent) {
				throw new Error('Cannot handle message: agent is not initialized')
			}
			console.log(' --- handle message --- ')
			state.history.push({
				type: 'message',
				payload: message
			})
			if (isVoiceAgent(agent)) {
				const reactAct = makeAction(sendAction, 'Agent', 'Streaming response')
				await reactAct('started')
				await agent.streamResponse()
				await reactAct('completed')
			} else {
				// text agent
				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 (e) {
						await agent.onFatalError(e as any)
					}
				}
			}
			addState()
		}
	}
}

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

function createState(schedule: DelayedActions, prev: string | null) {
	const state: AgentState = {
		id: new Date().toISOString(),
		name: new Date().toLocaleString(),
		history: [],
		schedule,
		prev,
		next: null
	}

	return state
}

function clone(obj: any) {
	return JSON.parse(JSON.stringify(obj))
}

function userMessage(st: AgentState): string | undefined {
	for (const entry of st.history) {
		if (entry.type === 'message') {
			const m = entry.payload as Message
			if (m.sender === 'user') {
				if (typeof (m as any).text === 'string' && (m as any).text.trim()) {
					return (m as any).text
				}
				if (typeof (m as any).imageUrl === 'string' && (m as any).imageUrl.trim()) {
					return `[image] ${(m as any).imageUrl}`
				}
				return JSON.stringify(m)
			}
		}
	}
	return undefined
}

async function loadSchemaFile(basePath: string, promptPath: string) {
	const dir = path.dirname(promptPath)
	const schemaPath = join(basePath, dir, 'schema.ts')
	console.log('loadSchemaFile', schemaPath)
	try {
		return await fs.promises.readFile(schemaPath, 'utf-8')
	} catch (e) {
		// console.error(e)
	}
}
