import WebSocket, {ErrorEvent, MessageEvent} from 'isomorphic-ws'

import {combineSignals} from '#common/async/aborts'
import {BlockingAtom} from '#common/async/blocking-atom'
import {BlockingQueue} from '#common/async/blocking-queue'
import {Breaker} from '#common/breaker'
import {normalizeError, ProtocolError} from '#common/errors'
import {z} from 'zod/v4'

import {parseJson} from './protocol'

/**
 * given a websocket, wait and take a single message off and return it
 *
 * @example
 * ```
 * const ws      = new WebSocket("wss://example.com/stream")
 * const timeout = timeoutSignal(5000)
 *
 * try {
 *   const msg = await takeSocket(ws, timeout.signal)
 *   doWhatever(msg)
 * } finally {
 *   timeout.cleanup()
 *   if (ws.readyState !== ws.CLOSED)
 *     ws.close();
 * }
 * ```
 */
export async function takeSocket(ws: WebSocket, signal?: AbortSignal): Promise<unknown> {
  signal?.throwIfAborted()

  const atom = new BlockingAtom()
  const breaker = new Breaker()

  const error = new AbortController()
  const multisignal = combineSignals(error.signal, signal)

  const onMessage = breaker.tripThen((m: MessageEvent) => {
    atom.set(m.data)
  })
  const onError = breaker.tripThen((e: unknown) => {
    error.abort(e)
  })
  const onClose = breaker.tripThen(() => {
    error.abort('closed')
  })

  try {
    ws.addEventListener('message', onMessage)
    ws.addEventListener('error', onError)
    ws.addEventListener('close', onClose)

    const data = await atom.get(multisignal)
    if (!data) {
      const cause = normalizeError(multisignal.reason)
      throw new ProtocolError('socket read aborted', 408, {cause})
    }

    return data
  } finally {
    ws.removeEventListener('message', onMessage)
    ws.removeEventListener('error', onError)
    ws.removeEventListener('close', onClose)
  }
}

/** exactly take socket, but will additionally apply a json decoding */
export async function takeSocketJson<T>(ws: WebSocket, schema: z.ZodType<T>, signal?: AbortSignal): Promise<T> {
  const data = await takeSocket(ws, signal)
  return parseJson.pipe(schema).parseAsync(data)
}

/** stream configuration options */

interface ConfigProps {
  maxDepth: number
  signal?: AbortSignal
}

export const STREAM_CONFIG_DEFAULT: ConfigProps = {
  maxDepth: 1000,
}

// symbols for iteration protocol

const yield$ = Symbol('yield$')
const error$ = Symbol('error$')
const end$ = Symbol('end$')

type StreamYield = [typeof yield$, unknown] | [typeof error$, Error] | [typeof end$]

/**
 * given a websocket, stream messages off the socket as an async generator
 *
 * @example
 * ```
 * const ws      = new WebSocket("wss://example.com/stream")
 * const timeout = timeoutSignal(5000)
 *
 * try {
 *   // signal will fire in 5s, so we'll take as many messages as we get until then
 *   for await (const msg of streamSocket(ws, { signal: timeout.signal })) {
 *     doWhatever(msg)
 *   }
 * } finally {
 *   timeout.cleanup()
 *   if (ws.readyState !== ws.CLOSED)
 *     ws.close();
 * }
 * ```
 */
export async function* streamSocket(ws: WebSocket, config_?: Partial<ConfigProps>) {
  const {signal, ...config} = {...STREAM_CONFIG_DEFAULT, ...(config_ || {})}
  signal?.throwIfAborted()

  // await incoming messages without blocking
  const queue = new BlockingQueue<StreamYield>(config.maxDepth)

  // if true, we're ignoring incoming messages until we drop the queue
  let inBackoffMode: boolean = false
  const backoffThresh = Math.floor(config.maxDepth * 0.9)

  // we don't want to keep processing after we've been closed
  const breaker = new Breaker()

  // define callback functions (need to be able to reference for removing them)

  const onMessage = breaker.untilTripped((m: MessageEvent) => {
    if (inBackoffMode) {
      console.warn('ignoring incoming message due to backpressure protection!')
      return
    }

    queue.enqueue([yield$, m.data])
    inBackoffMode = queue.depth > backoffThresh
    if (inBackoffMode) {
      console.warn('message stream will start dropping messages due to backpressure!')
    }
  })

  // todo: why are we getting this on client shutdown instead of onClose?
  const onError = breaker.tripThen((e: ErrorEvent) => {
    queue.enqueue([error$, normalizeError(e.error)])
  })

  const onClose = breaker.tripThen(() => {
    queue.enqueue([end$])
  })

  // finally get into our loop
  try {
    ws.addEventListener('message', onMessage)
    ws.addEventListener('error', onError)
    ws.addEventListener('close', onClose)

    while (true) {
      signal?.throwIfAborted()

      const [event, value] = await queue.dequeue(signal)

      // TODO: eslint thinks inBackoffMode is always falsey...
      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
      if (inBackoffMode && queue.depth < backoffThresh) {
        console.debug('message stream will stop dropping messages due to eased backpressure')
        inBackoffMode = false
      }

      switch (event) {
        case yield$:
          yield value
          continue
        case error$:
          throw value
        case end$:
          return
      }
    }
  } finally {
    ws.removeEventListener('message', onMessage)
    ws.removeEventListener('error', onError)
    ws.removeEventListener('close', onClose)
  }
}

/**
 * exactly stream socket, but will additionally apply a json decoding
 * messages not validating will end the stream with an error
 */
export async function* streamSocketJson(ws: WebSocket, config?: Partial<ConfigProps>): AsyncGenerator {
  for await (const message of streamSocket(ws, config)) {
    yield parseJson.parseAsync(message)
  }
}

/**
 * exactly stream socket, but will additionally apply a json decoding
 * messages not validating will end the stream with an error
 */
export async function* streamSocketSchema<T>(
  ws: WebSocket,
  schema: z.ZodType<T>,
  config?: Partial<ConfigProps>,
): AsyncGenerator<T> {
  const parser = parseJson.pipe(schema)
  for await (const message of streamSocket(ws, config)) {
    yield await parser.parseAsync(message)
  }
}