import { Transport, TransportType } from './Transport'
import { CallRequestDTO } from '../CallRequestDTO'
import { CallResponseDTO } from '../CallResponseDTO'
import { RPCClientIdentity } from '../RPCClientIdentity'
import { PromiseWrapper } from '../utils/PromiseWrapper'
import { sleep } from '../utils/sleep'
import { GetDebugLogger } from '../utils/debug-logger'
import { UnsubscribeCallback } from '../types'
import { RPCClientOptions, CallOptions } from '../RPCClient'
import { CallRequestTransportError } from './errors'
import { generateCorrelationId } from '../utils/generateCorrelationId'

const debug = GetDebugLogger('rpc:WebSocketTransport')

const RECONNECT_BASE_DELAY_MS = 1000
const RECONNECT_MAX_DELAY_MS = 30000
const IDENTITY_CONFIRMATION_TIMEOUT_MS = 2000

type ServerMessageHandler = (serverMessage: any) => void

interface WebSocketTransportOptions {
  host: string
  identity?: RPCClientIdentity
  onConnectionStatusChange?: (connected: boolean) => void
  rpcOptions: RPCClientOptions
}

export class WebSocketTransport implements Transport {
  private connectedToRemote: boolean = false
  private readonly host: string
  private identity?: RPCClientIdentity
  private identityConfirmationPromise?: Promise<void>
  private resolveIdentityConfirmation?: () => void
  private pendingPromisesForResponse: {
    [correlationId: string]: {
      promiseWrapper: PromiseWrapper<
        CallResponseDTO,
        CallResponseDTO | CallRequestTransportError
      >
      websocketId: number
    }
  } = {}
  private reconnectAttempt: number = 0
  private reconnectTimeoutId?: ReturnType<typeof setTimeout>
  private serverMessageHandlers: ServerMessageHandler[] = []
  private websocket?: WebSocket
  private websocketId?: number = undefined

  readonly name = 'WebSocketTransport'

  readonly type = TransportType.websocket

  constructor(readonly options: WebSocketTransportOptions) {
    debug('new WebSocketTransport()')
    this.host = options.host

    if (options.rpcOptions.transportOptions?.websocket?.autoConnect !== false) {
      this.connect()
    }
  }

  public isConnected = () => {
    return this.websocket?.readyState === WebSocket.OPEN
  }

  public connect = (): void => {
    this._connect()
  }

  public disconnect = (): void => {
    this._disconnect('public disconnect() called')
  }

  public sendClientMessageToServer = (msg: any) => {
    let strMsg = msg

    try {
      strMsg = JSON.stringify({ clientMessage: msg })
    } catch (err) {
      console.error(
        'WebSocketTransport.sendClientMessage() unable to stringify "msg"',
        msg,
      )
      return
    }

    debug('sendClientMessageToServer', strMsg)
    if (this.websocket?.readyState === WebSocket.OPEN) {
      this.websocket?.send(strMsg)
    } else {
      debug('WebSocket is not connected. Message not sent:', strMsg)
    }
  }

  public sendRequest = async (
    call: CallRequestDTO,
    opts: CallOptions,
  ): Promise<CallResponseDTO> => {
    debug('sendRequest', call)

    if (this.websocket?.readyState !== WebSocket.OPEN) {
      debug('sendRequest WebSocket is not connected')
      throw new CallRequestTransportError('WebSocket is not connected')
    }

    // Wait for identity confirmation using Promise instead of busy-wait
    if (this.identityConfirmationPromise) {
      debug('waiting for identity confirmation')
      await Promise.race([
        this.identityConfirmationPromise,
        new Promise<void>((_, reject) => {
          setTimeout(() => {
            reject(
              new CallRequestTransportError(
                'Timed out waiting for identity confirmation',
              ),
            )
          }, IDENTITY_CONFIRMATION_TIMEOUT_MS)
        }),
      ])
    }

    return this.sendCall(call, opts)
  }

  public setIdentity = async (identity?: RPCClientIdentity) => {
    debug('setIdentity', identity)

    if (!identity) {
      this.identity = undefined
      await this.resetConnection()

      return
    } else {
      this.identity = identity
    }

    if (!this.connectedToRemote) {
      debug('setIdentity is not connected to remote')
      return
    } else if (this.identityConfirmationPromise) {
      debug(
        'setIdentity returning early because identity confirmation is already in progress',
      )
      return
    }

    // Create a promise that other calls can await
    this.identityConfirmationPromise = new Promise<void>((resolve) => {
      this.resolveIdentityConfirmation = resolve
    })

    try {
      await this.sendCall({ identity }, {})
      debug('setIdentity with remote complete')
    } catch (e) {
      debug('setIdentity with remote error', e)
      throw e
    } finally {
      // Signal that identity confirmation is complete
      if (this.resolveIdentityConfirmation) {
        this.resolveIdentityConfirmation()
      }
      this.identityConfirmationPromise = undefined
      this.resolveIdentityConfirmation = undefined
    }
  }

  public subscribeToServerMessages = (
    handler: (msg: any) => void,
  ): UnsubscribeCallback => {
    this.serverMessageHandlers.push(handler)

    return () => {
      this.serverMessageHandlers = this.serverMessageHandlers.filter(
        (cb) => cb !== handler,
      )
    }
  }

  public unsubscribeFromServerMessages = (handler: (msg: any) => void) => {
    this.serverMessageHandlers = this.serverMessageHandlers.filter(
      (_handler) => _handler !== handler,
    )
  }

  public clearAllServerMessageHandlers = () => {
    this.serverMessageHandlers = []
  }

  public destroy = () => {
    debug('destroy')
    this._disconnect('destroy() called')
    this.clearAllServerMessageHandlers()
  }

  private _connect = () => {
    debug('connect', this.host)

    if (this.websocket) {
      debug('connect() returning early, websocket already exists')
      return
    }

    this.websocketId = Math.random()
    this.websocket = new WebSocket(this.host)

    const ws = this.websocket

    ws.onopen = () => {
      console.info(`[${new Date().toLocaleTimeString()}] WebSocket connected`)
      this.reconnectAttempt = 0 // Reset on successful connection
      this.setConnectedToRemote(true)
    }

    ws.onmessage = (msg) => {
      this.handleWebSocketMsg(msg)
    }

    ws.onclose = (e) => {
      if (this.connectedToRemote) {
        console.info(
          `[${new Date().toLocaleTimeString()}] WebSocket closed`,
          e.reason || e.code,
        )
      } else {
        debug('WebSocket closed, it was not connected to the remote')
      }

      this.setConnectedToRemote(false)

      // Clean up identity confirmation state
      if (this.resolveIdentityConfirmation) {
        this.resolveIdentityConfirmation()
      }
      this.identityConfirmationPromise = undefined
      this.resolveIdentityConfirmation = undefined

      this.websocket = undefined
      this.websocketId = undefined

      Object.entries(this.pendingPromisesForResponse).forEach(
        ([correlationId, { promiseWrapper }]) => {
          promiseWrapper.reject(
            new CallRequestTransportError('Websocket closed'),
          )
        },
      )

      this.pendingPromisesForResponse = {}

      if (!e.wasClean) {
        this.scheduleReconnect()
      }
    }

    ws.onerror = (err) => {
      if (this.connectedToRemote) {
        console.error('WebSocket encountered error: ', err)
      } else {
        debug('WebSocket errored, it was not connected to the remote')
      }

      ws.close()
    }
  }

  private _disconnect = async (reason?: string) => {
    debug('_disconnect')
    this.setConnectedToRemote(false)

    // Clean up identity confirmation state
    if (this.resolveIdentityConfirmation) {
      this.resolveIdentityConfirmation()
    }
    this.identityConfirmationPromise = undefined
    this.resolveIdentityConfirmation = undefined

    // Clear any pending reconnect
    if (this.reconnectTimeoutId) {
      clearTimeout(this.reconnectTimeoutId)
      this.reconnectTimeoutId = undefined
    }
    this.reconnectAttempt = 0

    this.websocket?.close(1000, reason)
    this.websocket = undefined
    this.websocketId = undefined
  }

  private handleServerMessage = (msg: { serverMessage: any }) => {
    this.serverMessageHandlers.forEach((handler) => {
      try {
        handler(msg.serverMessage)
      } catch (err) {
        console.warn(
          `WebSocketTransport.handleServerMessage() a serverMessageHandler errored when calling`,
          msg,
          'handler func:',
          handler,
        )
        console.error(err)
      }
    })
  }

  private handleWebSocketMsg = (msg: MessageEvent) => {
    debug('handleWebSocketMsg', msg)

    let json: any

    try {
      json = JSON.parse(msg.data)
    } catch (e: any) {
      debug('error parsing WS msg', e)
      return // Exit early if JSON parsing fails
    }

    // Check if json is defined and has the serverMessage property
    if (json && 'serverMessage' in json) {
      this.handleServerMessage(json)

      return
    }

    // Wrap CallResponseDTO construction in try/catch to handle malformed responses
    let callResponseDTO: CallResponseDTO
    try {
      callResponseDTO = new CallResponseDTO(json)
    } catch (e: any) {
      console.error(
        'RPCClient WebSocketTransport failed to construct CallResponseDTO from message',
        e,
        json,
      )
      return
    }

    if (!callResponseDTO.correlationId) {
      console.error(
        'RPCClient WebSocketTransport received unexpected msg from the server, not correlationId found in response.',
        json,
      )

      return
    }

    const pendingEntry =
      this.pendingPromisesForResponse[callResponseDTO.correlationId]

    if (pendingEntry) {
      // Validate that the response is from the current WebSocket connection
      // This prevents processing stale responses after a reconnect
      if (pendingEntry.websocketId === this.websocketId) {
        pendingEntry.promiseWrapper.resolve(callResponseDTO)
      } else {
        debug(
          'Ignoring response from stale WebSocket connection',
          callResponseDTO.correlationId,
        )
      }
    } else {
      console.warn(
        "rcvd WS msg/response that doesn't match any pending RPC's",
        json,
      )
    }
  }

  private resetConnection = async () => {
    debug('resetConnection')

    // Store current connection state
    const wasConnected = this.isConnected()

    await this._disconnect('WebSocketTransport#resetConnection')

    if (wasConnected) {
      await sleep(300) // delay to avoid race conditions
      await this._connect()
    }
    return this.isConnected()
  }

  private sendCall = async (
    call: CallRequestDTO,
    opts: CallOptions,
  ): Promise<CallResponseDTO> => {
    debug('sendCall', call)

    const correlationId = generateCorrelationId()

    call.correlationId = call.correlationId || correlationId

    // Timeout is handled by ClientManager.sendRequestWithTransport() for consistent
    // behavior across all transports (HTTP and WebSocket)
    const promiseWrapper = new PromiseWrapper<
      CallResponseDTO,
      CallResponseDTO | CallRequestTransportError
    >()

    // Check if websocket is connected before sending
    if (!this.isConnected()) {
      const error = new CallRequestTransportError(
        'Cannot send message: WebSocket is not connected',
      )
      debug('sendCall failed - websocket not connected', error)
      promiseWrapper.reject(error)
      return promiseWrapper.promise
    }

    this.websocket?.send(JSON.stringify(call))

    // Store with websocketId to validate responses are from current connection
    this.pendingPromisesForResponse[call.correlationId] = {
      promiseWrapper,
      websocketId: this.websocketId!,
    }

    return await promiseWrapper.promise.finally(() => {
      delete this.pendingPromisesForResponse[call.correlationId!]
    })
  }

  private setConnectedToRemote = (connected: boolean) => {
    debug(`setConnectedToRemote: ${connected}`)
    this.connectedToRemote = connected

    if (this.options.onConnectionStatusChange) {
      this.options.onConnectionStatusChange(connected)
    }

    if (connected && this.identity) {
      this.setIdentity(this.identity)
    }
  }

  private calculateReconnectDelay = (): number => {
    // Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s, 30s...
    const exponentialDelay =
      RECONNECT_BASE_DELAY_MS * Math.pow(2, this.reconnectAttempt)
    const cappedDelay = Math.min(exponentialDelay, RECONNECT_MAX_DELAY_MS)

    // Add jitter: +/- 20% to spread out reconnection attempts
    const jitter = cappedDelay * 0.2 * (Math.random() * 2 - 1)

    return Math.floor(cappedDelay + jitter)
  }

  private scheduleReconnect = (): void => {
    const delay = this.calculateReconnectDelay()
    this.reconnectAttempt++

    debug(`scheduling reconnect attempt ${this.reconnectAttempt} in ${delay}ms`)

    this.reconnectTimeoutId = setTimeout(() => {
      this.connect()
    }, delay)
  }
}