import type { BirpcOptions, PromisifyFn } from 'birpc'
import type { WebSocketEvents, WebSocketHandlers } from 'vitest'
import { createBirpc } from 'birpc'
import { parse, stringify } from 'flatted'
import { StateManager } from './state'

export interface VitestClientOptions {
  handlers?: Partial<WebSocketEvents>
  autoReconnect?: boolean
  reconnectInterval?: number
  reconnectTries?: number
  connectTimeout?: number
  reactive?: <T>(v: T, forKey: 'state' | 'idMap' | 'filesMap') => T
  ref?: <T>(v: T) => { value: T }
  WebSocketConstructor?: typeof WebSocket
}

export type VitestClientRpc = {
  [K in keyof WebSocketHandlers]: PromisifyFn<WebSocketHandlers[K]>
}

export interface VitestClient {
  ws: WebSocket
  state: StateManager
  rpc: VitestClientRpc
  // TODO: unused
  waitForConnection: () => Promise<void>
  reconnect: () => Promise<void>
}

export function createWsClient(url: string, options: VitestClientOptions = {}): VitestClient {
  const {
    handlers = {},
    autoReconnect = true,
    reconnectInterval = 2000,
    reconnectTries = 10,
    connectTimeout = 60000,
    reactive = v => v,
    WebSocketConstructor = globalThis.WebSocket,
  } = options

  let tries = reconnectTries
  let openPromise: Promise<void>
  const ctx = reactive<VitestClient>({
    ws: new WebSocketConstructor(url),
    state: new StateManager(),
    rpc: undefined!,
    waitForConnection: () => openPromise,
    reconnect,
  }, 'state')

  ctx.state.filesMap = reactive(ctx.state.filesMap, 'filesMap')
  ctx.state.idMap = reactive(ctx.state.idMap, 'idMap')

  let onMessage: (data: any) => void
  const functions: WebSocketEvents = {
    onTestAnnotate(testId, annotation) {
      handlers.onTestAnnotate?.(testId, annotation)
    },
    onTestArtifactRecord(testId, artifact) {
      handlers.onTestArtifactRecord?.(testId, artifact)
    },
    onSpecsCollected(specs, startTime) {
      specs?.forEach(([config, file]) => {
        ctx.state.clearFiles({ config }, [file])
      })
      handlers.onSpecsCollected?.(specs, startTime)
    },
    onCollected(files) {
      ctx.state.collectFiles(files)
      handlers.onCollected?.(files)
    },
    onTaskUpdate(packs, events) {
      ctx.state.updateTasks(packs)
      handlers.onTaskUpdate?.(packs, events)
    },
    onUserConsoleLog(log) {
      ctx.state.updateUserLog(log)
      handlers.onUserConsoleLog?.(log)
    },
    onFinished(files, errors, coverage, executionTime) {
      handlers.onFinished?.(files, errors, coverage, executionTime)
    },
    onFinishedReportCoverage() {
      handlers.onFinishedReportCoverage?.()
    },
  }

  const birpcHandlers = {
    post: msg => ctx.ws.send(msg),
    on: fn => (onMessage = fn),
    serialize: e =>
      stringify(e, (_, v) => {
        if (v instanceof Error) {
          return {
            name: v.name,
            message: v.message,
            stack: v.stack,
          }
        }
        return v
      }),
    deserialize: parse,
    timeout: -1,
  } satisfies BirpcOptions<WebSocketHandlers>

  ctx.rpc = createBirpc<WebSocketHandlers, WebSocketEvents>(
    functions,
    birpcHandlers,
  )

  async function reconnect(reset = false) {
    if (reset) {
      tries = reconnectTries
    }
    ctx.ws = new WebSocketConstructor(url)
    registerWS()
  }

  function registerWS() {
    openPromise = new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(
          new Error(
            `Cannot connect to the server in ${connectTimeout / 1000} seconds`,
          ),
        )
      }, connectTimeout)
      if (ctx.ws.OPEN === ctx.ws.readyState) {
        resolve()
      }
      // still have a listener even if it's already open to update tries
      ctx.ws.addEventListener('open', () => {
        tries = reconnectTries
        resolve()
        clearTimeout(timeout)
      })
    })
    ctx.ws.addEventListener('message', (v) => {
      onMessage(v.data)
    })
    ctx.ws.addEventListener('close', () => {
      tries -= 1
      if (autoReconnect && tries > 0) {
        setTimeout(reconnect, reconnectInterval)
      }
    })
  }

  registerWS()

  return ctx
}
