import { TASK_STATUS_ROUTE } from '../constants.js'
import type { RenegadeConfig } from '../createConfig.js'
import { AuthType } from '../utils/websocket.js'
import {
  type WebsocketWaiterParams,
  websocketWaiter,
} from '../utils/websocketWaiter.js'
import { getTaskHistory } from './getTaskHistory.js'
import type { WaitForTaskCompletionParameters } from './waitForTaskCompletion.js'

export async function waitForTaskCompletionWs(
  config: RenegadeConfig,
  parameters: WaitForTaskCompletionParameters,
): Promise<null | undefined> {
  const { id, timeout } = parameters
  const topic = TASK_STATUS_ROUTE(id)

  const wsWaiterParams: WebsocketWaiterParams = {
    config,
    topic,
    authType: AuthType.Wallet,
    messageHandler: (message: any) => {
      const parsedMessage = JSON.parse(message)
      if (
        parsedMessage.topic === topic &&
        parsedMessage.event.type === 'TaskStatusUpdate'
      ) {
        if (parsedMessage.event.status?.state === 'Completed') {
          return null
        }

        if (parsedMessage.event.status?.state === 'Failed') {
          throw new Error(`Task ${id} failed`)
        }
      }

      return undefined
    },
    prefetch: async () => {
      const taskHistory = await getTaskHistory(config)
      const task = taskHistory.get(id)

      if (task?.state === 'Completed') {
        return null
      }

      if (task?.state === 'Failed') {
        throw new Error(`Task ${id} failed`)
      }

      return undefined
    },
    timeout,
  }

  return websocketWaiter(wsWaiterParams)
}