Dynavera/src/stores/agentStore.ts

import { defineStore } from 'pinia'
import { ref } from 'vue'
import type { AgentEvent } from '../types/agent'

export const useAgentStore = defineStore('agent', () => {
    const isConnected = ref(false)
    const executionStatus = ref<'idle' | 'running' | 'completed' | 'failed'>('idle')
    const eventLog = ref<AgentEvent[]>([])
    const lastExecutionId = ref<string | null>(null)

    let socket: WebSocket | null = null

    const pushEvent = (evt: { type: string; message?: string; content?: unknown; timestamp?: string }) => {
        eventLog.value.unshift({
            type: evt.type,
            message: evt.message,
            content: evt.content,
            timestamp: evt.timestamp ? new Date(evt.timestamp) : new Date(),
        })
    }

    const connect = (agentId: string) => {
        if (socket) {
            socket.close()
            socket = null
        }

        const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
        const wsUrl = `${wsProtocol}://${window.location.host}/ws/mlstore/agents/${agentId}/`
        socket = new WebSocket(wsUrl)

        socket.onopen = () => {
            isConnected.value = true
            pushEvent({ type: 'connected', message: 'WebSocket connected' })
        }

        socket.onmessage = (event) => {
            try {
                const payload = JSON.parse(event.data)
                const type = payload.type

                if (type === 'connection') {
                    pushEvent({ type: 'connection', message: payload.message, content: payload })
                    return
                }

                if (type === 'execution_started') {
                    executionStatus.value = 'running'
                    lastExecutionId.value = payload.execution_id
                    pushEvent({ type: 'started', message: payload.message, content: payload })
                    return
                }

                if (type === 'mlstore_event') {
                    const eventType = payload.event_type || 'message'
                    pushEvent({
                        type: eventType,
                        content: payload.content,
                        timestamp: payload.timestamp,
                    })
                    return
                }

                if (type === 'execution_completed') {
                    executionStatus.value = 'completed'
                    pushEvent({ type: 'completed', content: payload.output_data, timestamp: payload.timestamp })
                    return
                }

                if (type === 'execution_error') {
                    executionStatus.value = 'failed'
                    pushEvent({ type: 'error', message: payload.error_message, content: payload })
                    return
                }

                pushEvent({ type: 'message', content: payload })
            } catch {
                pushEvent({ type: 'error', message: 'Invalid message received', content: event.data })
            }
        }

        socket.onerror = () => {
            pushEvent({ type: 'error', message: 'WebSocket error' })
        }

        socket.onclose = () => {
            isConnected.value = false
            pushEvent({ type: 'disconnected', message: 'WebSocket disconnected' })
        }
    }

    const disconnect = () => {
        if (socket) {
            socket.close()
            socket = null
        }
        isConnected.value = false
    }

    const startAgent = (data: {
        query?: string
        prompt?: string
        role_uuid?: string
        max_tokens?: number
        options?: Record<string, unknown>
    }) => {
        if (!socket || socket.readyState !== WebSocket.OPEN) return
        const prompt = data.query ?? data.prompt ?? ''
        const options = {
            ...(data.options ?? {}),
            ...(typeof data.max_tokens === 'number' ? { max_tokens: data.max_tokens } : {}),
        }
        socket.send(
            JSON.stringify({
                action: 'infer',
                input_data: {
                    prompt,
                    role_uuid: data.role_uuid,
                    options,
                },
            }),
        )
    }

    const startFineTune = (inputData?: Record<string, unknown>) => {
        if (!socket || socket.readyState !== WebSocket.OPEN) return
        socket.send(
            JSON.stringify({
                action: 'fine_tune',
                input_data: inputData ?? {},
            }),
        )
    }

    const stopAgent = (executionId?: string) => {
        if (!socket || socket.readyState !== WebSocket.OPEN) return
        socket.send(
            JSON.stringify({
                action: 'stop_agent',
                execution_id: executionId ?? lastExecutionId.value,
            }),
        )
    }

    const sendOnboardingProgress = (executionId: string, content: Record<string, unknown>) => {
        if (!socket || socket.readyState !== WebSocket.OPEN) return
        socket.send(
            JSON.stringify({
                action: 'onboarding_progress',
                execution_id: executionId,
                content,
            }),
        )
    }

    const resetLog = () => {
        eventLog.value = []
    }

    return {
        isConnected,
        executionStatus,
        eventLog,
        connect,
        disconnect,
        startAgent,
        startFineTune,
        stopAgent,
        sendOnboardingProgress,
        resetLog,
        lastExecutionId,
    }
})