import { ref, onUnmounted, type Ref } from 'vue'

export interface PeerStats {
  publicKey: string
  ip: string
  rxBytes: number
  txBytes: number
  lastHandshake: string
  isOnline: boolean
}

export interface TrafficDataPoint {
  timestamp: number
  rxBytes: number
  txBytes: number
}

export interface WebSocketStats {
  totalPeers: number
  activePeers: number
  totalRules: number
  trafficHistory: TrafficDataPoint[]
  peers: PeerStats[]
}

export interface UseWebSocketReturn {
  isConnected: Ref<boolean>
  stats: Ref<WebSocketStats>
  error: Ref<string | null>
  connect: () => void
  disconnect: () => void
}

const DEFAULT_STATS: WebSocketStats = {
  totalPeers: 0,
  activePeers: 0,
  totalRules: 0,
  trafficHistory: [],
  peers: []
}

export function useWebSocket(url: string = 'ws://localhost:8080/ws/stats'): UseWebSocketReturn {
  const isConnected = ref(false)
  const stats = ref<WebSocketStats>({ ...DEFAULT_STATS })
  const error = ref<string | null>(null)
  let ws: WebSocket | null = null
  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
  const reconnectDelay = 3000
  let trafficBuffer: TrafficDataPoint[] = []

  const connect = () => {
    if (ws && ws.readyState === WebSocket.OPEN) return

    error.value = null

    try {
      ws = new WebSocket(url)

      ws.onopen = () => {
        isConnected.value = true
        error.value = null
      }

      ws.onmessage = (event: MessageEvent) => {
        try {
          const data = JSON.parse(event.data)

          if (data.peers !== undefined) {
            stats.value.peers = data.peers
            stats.value.totalPeers = data.peers.length
            stats.value.activePeers = data.peers.filter((p: PeerStats) => p.isOnline).length
          }

          if (data.totalRules !== undefined) {
            stats.value.totalRules = data.totalRules
          }

          if (data.traffic) {
            const point: TrafficDataPoint = {
              timestamp: Date.now(),
              rxBytes: data.traffic.rxBytes || 0,
              txBytes: data.traffic.txBytes || 0
            }

            trafficBuffer.push(point)
            if (trafficBuffer.length > 60) trafficBuffer.shift()

            stats.value.trafficHistory = [...trafficBuffer]
          }

          if (data.stats) {
            stats.value = { ...stats.value, ...data.stats }
          }
        } catch (e) {
          console.error('Failed to parse WebSocket message:', e)
        }
      }

      ws.onerror = () => {
        error.value = 'WebSocket connection error'
      }

      ws.onclose = () => {
        isConnected.value = false
        ws = null
        scheduleReconnect()
      }
    } catch (e) {
      error.value = `Failed to connect: ${e}`
      scheduleReconnect()
    }
  }

  const scheduleReconnect = () => {
    if (reconnectTimer) clearTimeout(reconnectTimer)
    reconnectTimer = setTimeout(() => {
      connect()
    }, reconnectDelay)
  }

  const disconnect = () => {
    if (reconnectTimer) {
      clearTimeout(reconnectTimer)
      reconnectTimer = null
    }
    if (ws) {
      ws.close()
      ws = null
    }
    isConnected.value = false
  }

  onUnmounted(() => {
    disconnect()
  })

  return {
    isConnected,
    stats,
    error,
    connect,
    disconnect
  }
}
