import { ref, onUnmounted } from 'vue'
import { useChatStore } from '@/stores/chat.js'
import { useToast } from 'gnexus-ui-kit/vue'
import { t } from '@/i18n/index.js'

const WS_SCHEME = location.protocol === 'https:' ? 'wss' : 'ws'
const WS_BASE = import.meta.env.DEV
  ? `${WS_SCHEME}://${location.hostname}:8099`
  : `${WS_SCHEME}://${location.host}`

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

// Liveness. The server pings every 20 s (both mid-run and idle), so two
// missed intervals mean the socket is half-open: readyState still says OPEN,
// writes go into the void, and no close event ever arrives — the state a
// phone leaves behind after sleeping or switching networks.
const PING_INTERVAL_MS = 20000
const STALE_AFTER_MS = 45000
const WATCHDOG_INTERVAL_MS = 5000
const CONNECT_TIMEOUT_MS = 15000

function getWsUrl(sessionId) {
  const url = `${WS_BASE}/ws/sessions/${sessionId}`
  // SECURITY: token is stored in localStorage (XSS risk) and passed via query
  // param (visible in server access logs). Future: switch to auth message.
  const token = localStorage.getItem('navi_api_token')
  if (token) {
    return `${url}?api_token=${encodeURIComponent(token)}`
  }
  return url
}

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

  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
  let _isReconnect = false  // true when this open is a retry, not the first connect
  let _isConnecting = false // guard against redundant _connect() calls
  let _connectingSince = 0
  // Last frame received on the current socket — the only evidence that the
  // connection is still alive end to end.
  let lastMessageAt = 0
  let pingTimer = null
  let watchdogTimer = null
  let _listenersAttached = false
  let queuedSends = []

  function connect(id) {
    destroyed = false
    sessionId = id
    retryCount = 0
    reconnectFailed.value = false
    queuedSends = []
    _attachLifecycle()
    _connect()
  }

  function disconnect() {
    destroyed = true
    clearTimeout(retryTimer)
    _stopTimers()
    _detachLifecycle()
    ws?.close()
    ws = null
    _isConnecting = false
    queuedSends = []
    connected.value = false
    reconnecting.value = false
  }

  function send(payload) {
    // A socket that looks open but has been silent for two heartbeat windows
    // is a zombie: writing to it loses the message silently. Treat it exactly
    // like a closed socket — queue, reconnect, flush on open.
    if (ws?.readyState === WebSocket.OPEN && !_isStale()) {
      ws.send(JSON.stringify(payload))
      return true
    }
    if (queuedSends.length >= MAX_QUEUED_SENDS) queuedSends.shift()
    queuedSends.push(payload)
    reconnecting.value = true
    if (ws?.readyState === WebSocket.OPEN) {
      _forceReconnect()  // open but silent — a half-open socket, not a transport
    } else if (!ws || ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) {
      _connect()
    }
    return false
  }

  function _isStale() {
    return lastMessageAt !== 0 && Date.now() - lastMessageAt > STALE_AFTER_MS
  }

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

    ws.onopen = () => {
      _isConnecting = false
      connected.value = true
      reconnecting.value = false
      reconnectFailed.value = false
      _isReconnect = retryCount > 0
      retryCount = 0
      lastMessageAt = Date.now()
      _startTimers()
      // Whatever was mid-stream when the old socket died cannot be trusted:
      // the server replays a run that is still live from scratch, and answers
      // with session_sync when it is over. Dropping it here is what lets that
      // session_sync through — the reload is skipped while a stream is open.
      chat.dropStaleStream()
      _flushQueuedSends()
    }

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

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

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

  function _flushQueuedSends() {
    if (!ws || ws.readyState !== WebSocket.OPEN || !queuedSends.length) return
    const pending = queuedSends
    queuedSends = []
    for (const payload of pending) {
      ws.send(JSON.stringify(payload))
    }
  }

  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)
    }
  }

  // ── Liveness timers ──────────────────────────────────────────────────────

  function _startTimers() {
    _stopTimers()
    pingTimer = setInterval(() => {
      if (ws?.readyState === WebSocket.OPEN) {
        try { ws.send(JSON.stringify({ type: 'ping', ts: Date.now() })) } catch { /* closing */ }
      }
    }, PING_INTERVAL_MS)
    watchdogTimer = setInterval(_watchdog, WATCHDOG_INTERVAL_MS)
  }

  function _stopTimers() {
    clearInterval(pingTimer)
    clearInterval(watchdogTimer)
    pingTimer = null
    watchdogTimer = null
  }

  function _watchdog() {
    if (destroyed || !ws) return
    if (ws.readyState === WebSocket.CONNECTING) {
      // A handshake that never completes (captive portal, dead network) would
      // otherwise sit in CONNECTING forever — no close event, no retry.
      if (Date.now() - _connectingSince > CONNECT_TIMEOUT_MS) _forceReconnect()
      return
    }
    if (ws.readyState === WebSocket.OPEN && _isStale()) _forceReconnect()
  }

  // Drop the current socket and open a fresh one now, bypassing the backoff.
  function _forceReconnect() {
    if (destroyed) return
    const dead = ws
    ws = null
    _isConnecting = false
    if (dead) {
      dead.onopen = null
      dead.onmessage = null
      dead.onerror = null
      dead.onclose = null
      try { dead.close() } catch { /* already gone */ }
    }
    _stopTimers()
    retryCount = 0
    clearTimeout(retryTimer)
    connected.value = false
    reconnecting.value = true
    _connect()
  }

  // ── Lifecycle: resume from background, network back, bfcache ─────────────

  function _onVisibility() {
    if (destroyed || !sessionId) return
    if (document.visibilityState !== 'visible') return
    if (!ws || ws.readyState !== WebSocket.OPEN || _isStale()) {
      _forceReconnect()
      return
    }
    // Same socket, but another device (or a headless recall) may have advanced
    // the session while this tab was hidden.
    chat.checkForUpdates(sessionId)
  }

  function _onOnline() {
    if (destroyed || !sessionId) return
    // The transport that carried the old socket is gone even if the socket
    // does not know it yet.
    _forceReconnect()
  }

  function _onPageShow(event) {
    if (destroyed || !sessionId) return
    if (!event.persisted) return  // normal load — connect() already ran
    _forceReconnect()
  }

  function _attachLifecycle() {
    if (_listenersAttached || typeof window === 'undefined') return
    document.addEventListener('visibilitychange', _onVisibility)
    window.addEventListener('online', _onOnline)
    window.addEventListener('pageshow', _onPageShow)
    _listenersAttached = true
  }

  function _detachLifecycle() {
    if (!_listenersAttached) return
    document.removeEventListener('visibilitychange', _onVisibility)
    window.removeEventListener('online', _onOnline)
    window.removeEventListener('pageshow', _onPageShow)
    _listenersAttached = false
  }

  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 'planning_status':   chat.onPlanningStatus(event); break
      case 'plan_ready':        chat.onPlanReady(event); break
      case 'tool_started':      chat.onToolStarted(event); break
      case 'tool_call':         chat.onToolCall(event); break
      case 'terminal_output':    chat.onTerminalOutput(event); break
      case 'terminal_closed':  chat.onTerminalClosed(event); break
      case 'terminal_opened':  chat.onTerminalOpened(event); break
      // Background tasks: delivered out-of-band, outside the run's replay buffer
      case 'task_update': {
        // Toast the lifecycle transitions the user can't otherwise see while
        // working elsewhere in the chat. The store owns the session check and
        // the terminal-task suppression — its post-dispatch state is what
        // decides whether a toast fires.
        const id = event.task_id
        const existed = id ? Boolean(chat.backgroundTasks[id]) : false
        chat.onTaskUpdate(event)
        const now = id ? chat.backgroundTasks[id] : null
        const label = `${event.tool ?? t('common.task')} (${id})`
        if (event.status === 'running') {
          if (!existed && now) toast.info({ title: t('toast.bgTaskStarted'), text: label })
        } else if (existed) {
          const done = { title: t('toast.bgTaskFinished'), text: `${label} — ${event.status}` }
          if (event.status === 'completed') toast.success(done)
          else toast.error(done)
        }
        break
      }
      case 'message_queued':    chat.onMessageQueued(event); break
      case 'stream_delta':      chat.onStreamDelta(event.delta ?? ''); break
      case 'stream_end':        chat.onStreamEnd(event); break
      case 'stream_stopped':    chat.onStreamStopped(); break
      // ui_component is delivered via tool_call metadata, not a separate WS event.
      case 'profile_switched':  chat.onProfileSwitched(event); break
      case 'model_info':        chat.onModelInfo(event); break
      case 'compression_started':chat.onCompressionStarted(event); break
      case 'context_compressed':chat.onContextCompressed(event); break
      case 'error':             chat.onError(event); break
      case 'replay_start':      chat.onReplayStart(); break
      case 'replay_end':        chat.onReplayEnd(); break
      case 'heartbeat':         break  // keep-alive ping, no action needed
      case 'pong':              break  // answer to our own liveness probe
      case 'recall_update':
        chat.onRecallUpdate(event)
        break
      case 'session_sync':
        // Reload session history whenever the server signals persistence.
        // This handles both reconnect-after-finish and edge cases where
        // the client missed the final stream_end event.
        // Skip while a stream is active — reloadSession would orphan the
        // live streaming message and break delta/tool rendering. A stream
        // that died with its socket was already dropped on reconnect, so a
        // finished run always gets through here.
        if (sessionId && !chat.streaming) {
          chat.reloadSession(sessionId)
        }
        break
      case 'mcp_status_update':
        if (event.status === 'connected') {
          toast.success({
            title: t('toast.mcpConnected'),
            text: `${event.server_name} (${t('common.toolsCount', { n: event.tool_count ?? '?' })})`,
          })
        } else {
          toast.error({
            title: t('toast.mcpDisconnected'),
            text: event.server_name,
          })
        }
        break
    }
  }

  onUnmounted(disconnect)

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