import type {
	DelayedActions,
	LogEntry,
	AgentPayload,
	AgentSateHeader,
	WsMessage
} from '../lib/ws-message'
import { type Action, type Message } 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)
		}
		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 })
	}

	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 }
		})
	}
}
