import { ref, onUnmounted } from 'vue'
import { useChatStore } from '@/stores/chat.js'

const WS_BASE = import.meta.env.DEV
  ? `ws://${location.hostname}:8000`
  : `ws://${location.host}`

const MAX_FAST_RETRIES = 3
const FAST_RETRY_BASE_MS = 1000
const BACKGROUND_RETRY_MS = 15000

// Auth extension point — wire in token here when multi-user is ready
function getWsUrl(sessionId) {
  return `${WS_BASE}/ws/sessions/${sessionId}`
}

export function useWebSocket() {
  const chat = useChatStore()

  const connected = ref(false)
  const reconnecting = ref(false)
  const reconnectFailed = ref(false)  // true after MAX_FAST_RETRIES exhausted

  let ws = null
  let sessionId = null
  let retryCount = 0
  let retryTimer = null
  let destroyed = false

  function connect(id) {
    sessionId = id
    retryCount = 0
    reconnectFailed.value = false
    _connect()
  }

  function disconnect() {
    destroyed = true
    clearTimeout(retryTimer)
    ws?.close()
    ws = null
    connected.value = false
    reconnecting.value = false
  }

  function send(payload) {
    if (ws?.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(payload))
    }
  }

  function _connect() {
    if (destroyed) return
    if (ws) {
      ws.onclose = null  // Prevent old socket's onclose from triggering reconnect loop
      ws.close()
    }
    ws = new WebSocket(getWsUrl(sessionId))

    ws.onopen = () => {
      connected.value = true
      reconnecting.value = false
      reconnectFailed.value = false
      retryCount = 0
    }

    ws.onmessage = (e) => {
      let event
      try { event = JSON.parse(e.data) } catch { return }
      _dispatch(event)
    }

    ws.onclose = () => {
      connected.value = false
      if (!destroyed) _scheduleReconnect()
    }

    ws.onerror = () => {
      // onclose fires after onerror, no extra handling needed
    }
  }

  function _scheduleReconnect() {
    if (destroyed) return
    reconnecting.value = true
    retryCount++

    if (retryCount <= MAX_FAST_RETRIES) {
      const delay = FAST_RETRY_BASE_MS * Math.pow(2, retryCount - 1)
      retryTimer = setTimeout(_connect, delay)
    } else {
      // Switch to background retry
      reconnectFailed.value = true
      reconnecting.value = false
      retryTimer = setTimeout(() => {
        if (!destroyed) {
          reconnecting.value = true
          _connect()
        }
      }, BACKGROUND_RETRY_MS)
    }
  }

  function _dispatch(event) {
    switch (event.type) {
      case 'stream_start':      chat.onStreamStart(); break
      case 'thinking_delta':    chat.onThinkingDelta(event.delta ?? ''); break
      case 'thinking_end':      chat.onThinkingEnd(); break
      case 'turn_thinking':     chat.onTurnThinking(event); break
      case 'plan_ready':        chat.onPlanReady(event); break
      case 'tool_started':      chat.onToolStarted(event); break
      case 'tool_call':         chat.onToolCall(event); break
      case 'stream_delta':      chat.onStreamDelta(event.delta ?? ''); break
      case 'stream_end':        chat.onStreamEnd(event); break
      case 'stream_stopped':    chat.onStreamStopped(); break
      case 'profile_switched':  chat.onProfileSwitched(event); break
      case 'context_compressed':chat.onContextCompressed(event); break
      case 'error':             chat.onError(event); break
    }
  }

  onUnmounted(disconnect)

  return { connected, reconnecting, reconnectFailed, connect, disconnect, send }
}