import type { StreamEvent } from '../types'

type EventCallback = (event: StreamEvent) => void

export class ChatWebSocket {
  private ws: WebSocket | null = null
  private url: string
  private reconnectAttempts = 0
  private maxReconnectAttempts = 5
  private heartbeatInterval: ReturnType<typeof setInterval> | null = null
  private listeners: Set<EventCallback> = new Set()
  private _connected = false

  constructor(url?: string) {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    this.url = url || `${protocol}//${window.location.host}/api/ws/chat`
  }

  get connected(): boolean {
    return this._connected
  }

  connect(): void {
    if (this.ws?.readyState === WebSocket.OPEN) return

    this.ws = new WebSocket(this.url)

    this.ws.onopen = () => {
      this._connected = true
      this.reconnectAttempts = 0
      this.startHeartbeat()
    }

    this.ws.onmessage = (event) => {
      try {
        const data: StreamEvent = JSON.parse(event.data)
        this.listeners.forEach(cb => cb(data))
      } catch { /* ignore parse errors */ }
    }

    this.ws.onclose = () => {
      this._connected = false
      this.stopHeartbeat()
      this.tryReconnect()
    }

    this.ws.onerror = () => {
      this._connected = false
    }
  }

  disconnect(): void {
    this.maxReconnectAttempts = 0 // Prevent reconnect
    this.stopHeartbeat()
    this.ws?.close()
    this.ws = null
    this._connected = false
  }

  send(question: string, conversationId?: string): void {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      this.connect()
      // Queue message for after connection
      setTimeout(() => this.send(question, conversationId), 500)
      return
    }

    this.ws.send(JSON.stringify({
      question,
      conversation_id: conversationId,
    }))
  }

  sendRaw(data: object): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(data))
    }
  }

  onEvent(callback: EventCallback): () => void {
    this.listeners.add(callback)
    return () => this.listeners.delete(callback)
  }

  private startHeartbeat(): void {
    this.heartbeatInterval = setInterval(() => {
      if (this.ws?.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ type: 'ping' }))
      }
    }, 30000)
  }

  private stopHeartbeat(): void {
    if (this.heartbeatInterval) {
      clearInterval(this.heartbeatInterval)
      this.heartbeatInterval = null
    }
  }

  private tryReconnect(): void {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) return
    this.reconnectAttempts++
    const delay = Math.min(1000 * Math.pow(2, this.reconnectAttempts), 30000)
    setTimeout(() => this.connect(), delay)
  }
}

export const chatWs = new ChatWebSocket()
