import { ref, onUnmounted } from 'vue'
import { useChatStore } from '@/stores/chat.js'
import { useToast } from 'gnexus-ui-kit/vue'
import { t } from '@/i18n/index.js'
import { setWsSend } from '@/composables/wsBridge.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

// How long a written message may wait for the server's ack. Writing to an open
// socket proves only that the browser accepted the bytes — the server may have
// dropped the frame, and a dead peer's close event can arrive much later.
const ACK_TIMEOUT_MS = 15000

// Application close codes the server uses for a socket that must not be
// retried: the session is gone (4004) or the user may not open it (4003 —
// anonymous, access denied, admin-only profile).
const FATAL_CLOSE_CODES = new Set([4003, 4004])

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
  // { code, reason } once the server closed with a code that retrying cannot
  // fix. The UI shows why and offers the one action that helps (sign in, or
  // leave the session).
  const fatalClose = ref(null)

  let ws = null
  let wsSessionId = null  // the session the current socket was opened for
  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 = []
  // client_id → timeout that marks the bubble failed if no ack arrives.
  let ackTimers = new Map()
  // Set when the server closed the socket with a code retrying cannot fix.
  let _fatal = false

  function connect(id) {
    destroyed = false
    // A socket opened for another session must not survive the switch: its
    // events would land in the new transcript and its replay would stream the
    // previous session's answer into it. This covers the handshake too — a
    // CONNECTING socket for A is closed here instead of being allowed to open
    // late, after the client has already moved on to B.
    if (ws && wsSessionId !== id) {
      _detachSocket()
      _stopTimers()
      _clearAckTimers()
      connected.value = false
    }
    // A new connect (another session, or an explicit retry) is a fresh start:
    // a fatal close belonged to the session the user just left.
    _fatal = false
    fatalClose.value = null
    sessionId = id
    retryCount = 0
    // A backoff timer scheduled by the previous socket would otherwise fire
    // after this connect and open a second socket for the same session.
    clearTimeout(retryTimer)
    retryTimer = null
    reconnectFailed.value = false
    queuedSends = []
    _attachLifecycle()
    _connect()
  }

  function disconnect() {
    destroyed = true
    // The socket is deliberately gone — a fatal close's explanation must not
    // outlive the session view that caused it.
    _fatal = false
    fatalClose.value = null
    clearTimeout(retryTimer)
    _stopTimers()
    _detachLifecycle()
    _detachSocket()
    _clearAckTimers()
    queuedSends = []
    connected.value = false
    reconnecting.value = false
  }

  /**
   * Write a payload to the server, or park it until the socket is back.
   *
   * Returns what actually happened, so a caller can tell the user instead of
   * losing a message silently:
   *   'sent'     — written to an open, live socket;
   *   'queued'   — the socket is down/stale; it is flushed on reconnect;
   *   'rejected' — the queue is full; this payload was NOT kept.
   *
   * The queue keeps its oldest entries on overflow. Dropping the oldest to
   * make room is what silently lost a message the user had already been shown
   * as sent ("queued", from the bubble).
   *
   * A payload with a `client_id` also gets a delivery deadline: the server
   * answers `message_accepted` (or `message_queued`) with that id, and silence
   * past ACK_TIMEOUT_MS marks the bubble as not delivered — see _write/_armAck.
   */
  function send(payload) {
    // A retry cannot fix a session that is gone or not ours; parking the
    // payload in a queue nobody will flush would hide that from the user.
    if (_fatal) return 'rejected'
    // 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 && wsSessionId === sessionId && !_isStale()) {
      _write(payload)
      return 'sent'
    }
    if (queuedSends.length >= MAX_QUEUED_SENDS) {
      _ensureReconnect()
      return 'rejected'
    }
    queuedSends.push(payload)
    _ensureReconnect()
    return 'queued'
  }

  // Write one payload and, for a message that asked for an ack, start its
  // deadline. Anything else (ping, form submit) is fire-and-forget.
  function _write(payload) {
    ws.send(JSON.stringify(payload))
    if (payload?.type !== 'message' || !payload.client_id) return
    _armAck(payload.client_id)
  }

  function _armAck(clientId) {
    _clearAck(clientId)
    ackTimers.set(
      clientId,
      setTimeout(() => {
        ackTimers.delete(clientId)
        // The server's read loop does not take a frame sent during an active
        // run until that run ends: the message is parked there, not lost. Do
        // not accuse it — wait for the turn to finish and give it a new window.
        if (chat.streaming) {
          _armAck(clientId)
          return
        }
        chat.markMessageFailed(clientId, 'unconfirmed')
      }, ACK_TIMEOUT_MS),
    )
  }

  function _clearAck(clientId) {
    if (!clientId) return
    const timer = ackTimers.get(clientId)
    if (timer) clearTimeout(timer)
    ackTimers.delete(clientId)
  }

  function _clearAckTimers() {
    for (const timer of ackTimers.values()) clearTimeout(timer)
    ackTimers.clear()
  }

  function _ensureReconnect() {
    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()
    }
  }

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

  // Detach every handler before closing: a socket we are done with must not be
  // able to run a callback (an old close scheduling a reconnect, or a late
  // message landing in the transcript of the session we just opened).
  function _detachSocket() {
    const old = ws
    ws = null
    wsSessionId = null
    _isConnecting = false
    if (!old) return
    old.onopen = null
    old.onmessage = null
    old.onerror = null
    old.onclose = null
    try { old.close() } catch { /* already gone */ }
  }

  function _connect() {
    if (destroyed || _fatal) return
    // A handshake already in flight for this session is exactly what we want;
    // opening a second socket would leave two transports into one transcript.
    if (_isConnecting && ws && wsSessionId === sessionId) return
    if (ws) _detachSocket()
    _isConnecting = true
    _connectingSince = Date.now()
    const target = sessionId
    const socket = new WebSocket(getWsUrl(target))
    ws = socket
    wsSessionId = target

    socket.onopen = () => {
      if (ws !== socket) return  // replaced while the handshake was in flight
      _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()
    }

    socket.onmessage = (e) => {
      // A frame from a socket that is no longer the current one (or from a
      // session the client has left) must never touch the transcript.
      if (ws !== socket || wsSessionId !== sessionId) return
      lastMessageAt = Date.now()
      let event
      try { event = JSON.parse(e.data) } catch { return }
      _dispatch(event)
    }

    socket.onclose = (event) => {
      if (ws !== socket) return
      _isConnecting = false
      connected.value = false
      _stopTimers()
      const code = event?.code
      if (FATAL_CLOSE_CODES.has(code)) {
        // Retrying a session that does not exist or is not ours loops forever
        // against the same refusal. Stop, and let the UI say what happened.
        _fatal = true
        reconnecting.value = false
        reconnectFailed.value = false
        clearTimeout(retryTimer)
        retryTimer = null
        fatalClose.value = { code, reason: event?.reason || '' }
        return
      }
      if (!destroyed) _scheduleReconnect()
    }

    socket.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) {
      // The deadline starts now, not when the payload was parked: until this
      // write there was nothing for the server to acknowledge.
      _write(payload)
    }
  }

  function _scheduleReconnect() {
    if (destroyed || _fatal) 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 || _fatal || !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 || _fatal) return
    _detachSocket()
    _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
      }
      // The server has taken ownership of a message this tab sent. Its
      // delivery deadline ends here; the run itself is next.
      case 'message_accepted':  _clearAck(event.client_id); break
      case 'message_queued':
        _clearAck(event.client_id)
        chat.onMessageQueued(event)
        break
      // Queue overflow dropped the oldest message(s) — they will never run, so
      // their bubbles stop looking queued.
      case 'message_dropped':   chat.onMessageDropped(event); break
      case 'stream_delta':      chat.onStreamDelta(event.delta ?? ''); break
      case 'stream_end':        chat.onStreamEnd(event); break
      case 'stream_stopped':    chat.onStreamStopped(event); 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 'skill_proposal_decided':
        chat.onProposalDecided(event)
        break
      case 'session_sync':
        // The server sends this on every (re)connect, and again when a run
        // reattached to a dead socket finishes. Reloading history here made
        // every session open fetch the same page twice (~1 MB of JSON each),
        // because the initial load had just populated it. The cheap check —
        // one /meta request — reloads only when the message count actually
        // moved. Skipped while a load or a stream owns the transcript.
        if (sessionId && !chat.streaming && !chat.loading) {
          chat.checkForUpdates(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
    }
  }

  // Cards in the message tree reach the socket through this bridge rather than
  // an inject() chain — see wsBridge.js for why.
  setWsSend(send)
  onUnmounted(() => {
    setWsSend(null)
    disconnect()
  })

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