import type {
	DelayedActions,
	LogEntry,
	AgentPayload,
	AgentSateHeader,
	WsMessage
} from '../../lib/ws-message'
import {
	type Action,
	type ApiCallTrace,
	type ConversationalTrace,
	type Message,
	type StepTrace
} from '@recombine-ai/engine'
import { WebSocket } from 'ws'

export class MessageDispatcher {
	constructor(private ws: WebSocket) {}

	private sendWsMessage(message: WsMessage) {
		if (message.channel === 'app-events') {
			console.log('<-- ', message.channel, message.event.type)
		} else if (message.channel === 'chat') {
			console.log('<-- ', message.channel, message.message)
		} else if (message.channel === 'step-trace') {
			console.log('<-- ', message.channel, message.trace?.name)
		} else if (message.channel === 'api-call-trace') {
			console.log(
				'<-- ',
				message.channel,
				`${message.trace?.operationName} (${message.trace?.method} ${message.trace?.path})`
			)
		} else if (message.channel === 'conversational-trace') {
			console.log('<-- ', message.channel, `${message.trace?.eventName} (${message.trace?.role})`)
		}
		this.ws.send(JSON.stringify(message))
	}

	sendAction = (action: Action) => {
		this.sendWsMessage({ channel: 'actions', action })
	}

	sendMessage = async (message: Message) => {
		this.sendWsMessage({ channel: 'chat', message })
	}

	sendLog = async (message: LogEntry) => {
		this.sendWsMessage({ channel: 'log', message })
	}

	sendTrace = async (trace: StepTrace) => {
		this.sendWsMessage({ channel: 'step-trace', trace })
	}

	sendApiCallTrace = async (trace: ApiCallTrace) => {
		this.sendWsMessage({ channel: 'api-call-trace', trace })
	}

	sendConversationalTrace = async (trace: ConversationalTrace) => {
		this.sendWsMessage({ channel: 'conversational-trace', trace })
	}

	sendSchedule = async (actions: DelayedActions) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'delayed-actions-updated', actions }
		})
	}

	sendAgentSelected = async (agent: AgentPayload) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'agent-selected', agent }
		})
	}

	sendContext = async (context: Record<string, any>, promptContext: Record<string, any>) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'context-updated', context, promptContext }
		})
	}

	sendPrompt = async (prompt: string) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'prompt-updated', prompt }
		})
	}

	sendLoaded = async (agents: string[]) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'agents-loaded', agents }
		})
	}

	sendUpdateStates = async (states: AgentSateHeader[]) => {
		this.sendWsMessage({
			channel: 'app-events',
			event: { type: 'update-states', states }
		})
	}
}
