import { defineStore } from 'pinia'
import { ref } from 'vue'
import { BACKOFF_BASE_MS, BACKOFF_MAX_MS, BACKOFF_MAX_ATTEMPTS } from './agentBackoff'

export type KaResponse = {
    answer: string
    updatedPage: boolean
    revisedPageBody: string | null
}

type PendingRequest = {
    resolve: (value: KaResponse) => void
    reject: (reason: string) => void
}

export const useKaStore = defineStore('ka', () => {
    const isConnected = ref(false)
    const isAsking = ref(false)
    const streamBuffer = ref('')
    const socket = ref<WebSocket | null>(null)

    let currentUrl = ''
    let reconnectAttempts = 0
    let reconnectTimer: ReturnType<typeof setTimeout> | null = null
    let intentionalClose = false
    let pending: PendingRequest | null = null

    const clearReconnectTimer = () => {
        if (reconnectTimer !== null) {
            clearTimeout(reconnectTimer)
            reconnectTimer = null
        }
    }

    const scheduleReconnect = () => {
        if (reconnectAttempts >= BACKOFF_MAX_ATTEMPTS) return
        const delay = Math.min(BACKOFF_BASE_MS * 2 ** reconnectAttempts, BACKOFF_MAX_MS)
        reconnectAttempts++
        reconnectTimer = setTimeout(() => {
            if (!intentionalClose) openSocket(currentUrl)
        }, delay)
    }

    const openSocket = (url: string) => {
        socket.value = new WebSocket(url)

        socket.value.onopen = () => {
            reconnectAttempts = 0
            isConnected.value = true
        }

        socket.value.onmessage = (event) => {
            try {
                const payload = JSON.parse(event.data)
                if (payload.type === 'stream_chunk') {
                    streamBuffer.value += payload.message || ''
                } else if (payload.type === 'completed') {
                    isAsking.value = false
                    streamBuffer.value = ''
                    const content = payload.content || {}
                    pending?.resolve({
                        answer: payload.message || '',
                        updatedPage: Boolean(content.updated_page),
                        revisedPageBody: content.revised_page_body ?? null,
                    })
                    pending = null
                } else if (payload.type === 'error') {
                    isAsking.value = false
                    streamBuffer.value = ''
                    pending?.reject(payload.message || 'KA error')
                    pending = null
                }
            } catch (e) {
                console.error('KA store message error', e)
            }
        }

        socket.value.onclose = (event) => {
            isConnected.value = false
            if (!intentionalClose && event.code !== 1000) {
                scheduleReconnect()
            }
        }
    }

    const connect = (sessionUuid: string) => {
        intentionalClose = false
        clearReconnectTimer()
        reconnectAttempts = 0

        if (socket.value) {
            socket.value.close()
            socket.value = null
        }

        const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
        currentUrl = `${wsProtocol}://${window.location.host}/ws/onboarding/knowledge/${sessionUuid}/`
        openSocket(currentUrl)
    }

    const disconnect = () => {
        intentionalClose = true
        clearReconnectTimer()
        if (pending) {
            pending.reject('Disconnected')
            pending = null
        }
        isAsking.value = false
        if (socket.value) {
            socket.value.close()
            socket.value = null
        }
        isConnected.value = false
    }

    const ask = (
        pageUuid: string,
        userMessage: string,
        mode: 'separate' | 'update_page',
    ): Promise<KaResponse> => {
        return new Promise((resolve, reject) => {
            if (!socket.value || socket.value.readyState !== WebSocket.OPEN) {
                return reject('Not connected')
            }
            isAsking.value = true
            streamBuffer.value = ''
            pending = { resolve, reject }
            socket.value.send(
                JSON.stringify({ action: 'ask', page_uuid: pageUuid, message: userMessage, mode }),
            )
        })
    }

    return { isConnected, isAsking, streamBuffer, connect, disconnect, ask }
})
