'use client'

import {
  createContext,
  createElement,
  useCallback,
  useContext,
  useEffect,
  useRef,
  useState,
  type ReactNode,
} from 'react'
import {
  type AvatarPlaybackEvent,
  type ChatMessage,
  type LobsterStatusMessage,
  type RelayMessage,
  type SessionEventMessage,
  HEARTBEAT_INTERVAL_MS,
  parseRelayMessage,
  toAvatarPlaybackEvent,
} from '@/lib/protocol'
import { useAuth } from './useAuth'
import { updateOnlineLobsters } from './useLobsters'

export type ConnectionState = 'connecting' | 'connected' | 'disconnected' | 'error'

const RELAY_URL = process.env.NEXT_PUBLIC_RELAY_URL ?? 'http://localhost:8787'

// ── Session handler types ────────────────────────────────────────

export interface SessionHandlers {
  onMessage?: (msg: ChatMessage) => void
  onTyping?: (isTyping: boolean) => void
  onSessionEvent?: (msg: SessionEventMessage) => void
  onHistoryResponse?: (messages: ChatMessage[], hasMore: boolean) => void
  onHistoryClearResult?: (ok: boolean, error?: string) => void
  onLobsterStatus?: (
    action: LobsterStatusMessage['action'],
    lobsters: LobsterStatusMessage['lobsters'],
  ) => void
}

interface SessionRegistration {
  handlers: React.RefObject<SessionHandlers>
}

// ── Context type ─────────────────────────────────────────────────

interface RelayContextValue {
  state: ConnectionState
  avatarEvent: (AvatarPlaybackEvent & { id: string }) | null
  publishAvatarEvent: (event: AvatarPlaybackEvent) => void
  send: (msg: RelayMessage) => void
  joinSession: (sessionId: string, handlers: React.RefObject<SessionHandlers>) => void
  leaveSession: (sessionId: string) => void
}

const RelayContext = createContext<RelayContextValue | null>(null)

// ── Provider ─────────────────────────────────────────────────────

export function RelayProvider({ children }: { children: ReactNode }) {
  const { token, signOut } = useAuth()
  const [state, setState] = useState<ConnectionState>('disconnected')
  const [avatarEvent, setAvatarEvent] = useState<(AvatarPlaybackEvent & { id: string }) | null>(null)

  const wsRef = useRef<WebSocket | null>(null)
  const heartbeatRef = useRef<ReturnType<typeof setInterval> | null>(null)
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const reconnectAttemptsRef = useRef(0)
  const invalidSessionHandledRef = useRef(false)
  const maxReconnect = 10

  /** sessionId → registration mapping for dynamic session subscriptions */
  const sessionHandlersRef = useRef(new Map<string, SessionRegistration>())

  const cleanup = useCallback(() => {
    if (heartbeatRef.current) {
      clearInterval(heartbeatRef.current)
      heartbeatRef.current = null
    }
    if (reconnectTimerRef.current) {
      clearTimeout(reconnectTimerRef.current)
      reconnectTimerRef.current = null
    }
  }, [])

  const send = useCallback((msg: RelayMessage) => {
    const ws = wsRef.current
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(msg))
    }
  }, [])

  const publishAvatarEvent = useCallback((event: AvatarPlaybackEvent) => {
    setAvatarEvent({
      ...event,
      id: `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
    })
  }, [])

  const connect = useCallback(() => {
    if (!token) return
    cleanup()
    setState('connecting')

    const wsUrl = `${RELAY_URL.replace(/^http/, 'ws')}/ws/__global__`
    const ws = new WebSocket(wsUrl)
    wsRef.current = ws

    ws.onopen = () => {
      const authMessage: RelayMessage = {
        type: 'auth',
        role: 'browser',
        token,
        sessionId: '__global__',
      }
      ws.send(JSON.stringify(authMessage))
      publishAvatarEvent({ type: 'auth' })
    }

    ws.onmessage = (event) => {
      const msg = parseRelayMessage(event.data as string)
      if (!msg) return

      switch (msg.type) {
        case 'auth_result': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          if (msg.ok) {
            setState('connected')
            reconnectAttemptsRef.current = 0
            invalidSessionHandledRef.current = false
            // Start heartbeat
            heartbeatRef.current = setInterval(() => {
              if (ws.readyState === WebSocket.OPEN) {
                ws.send(JSON.stringify({ type: 'ping', timestamp: Date.now() }))
              }
            }, HEARTBEAT_INTERVAL_MS)

            // Re-join any sessions that were registered before reconnect
            for (const sessionId of sessionHandlersRef.current.keys()) {
              ws.send(JSON.stringify({ type: 'session_join', sessionId }))
            }
          } else {
            setState('error')
            if (
              msg.error === 'Invalid session token' &&
              !invalidSessionHandledRef.current
            ) {
              invalidSessionHandledRef.current = true
              cleanup()
              void signOut()
            }
            ws.close(4001, 'Auth failed')
          }
          break
        }

        case 'lobster_status': {
          const lsm = msg as LobsterStatusMessage
          publishAvatarEvent(toAvatarPlaybackEvent(lsm))
          // Update shared lobster state
          updateOnlineLobsters(lsm.action, lsm.lobsters)
          // Also dispatch to any session-level handler that cares
          for (const reg of sessionHandlersRef.current.values()) {
            reg.handlers.current?.onLobsterStatus?.(lsm.action, lsm.lobsters)
          }
          break
        }

        case 'message': {
          const chatMsg = msg as ChatMessage
          publishAvatarEvent(toAvatarPlaybackEvent(chatMsg))
          const reg = sessionHandlersRef.current.get(chatMsg.sessionId)
          reg?.handlers.current?.onMessage?.(chatMsg)
          break
        }

        case 'typing': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          if (msg.from === 'ai') {
            const reg = sessionHandlersRef.current.get(msg.sessionId)
            reg?.handlers.current?.onTyping?.(msg.isTyping)
          }
          break
        }

        case 'session_event': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          const reg = sessionHandlersRef.current.get(msg.sessionId)
          reg?.handlers.current?.onSessionEvent?.(msg)
          break
        }

        case 'history_response': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          const reg = sessionHandlersRef.current.get(msg.sessionId)
          reg?.handlers.current?.onHistoryResponse?.(msg.messages, msg.hasMore)
          break
        }

        case 'history_clear_result': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          if ('sessionId' in msg) {
            const reg = sessionHandlersRef.current.get(msg.sessionId)
            reg?.handlers.current?.onHistoryClearResult?.(msg.ok, msg.error)
          }
          break
        }

        case 'ack':
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          break

        case 'pong':
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          break

        case 'error': {
          publishAvatarEvent(toAvatarPlaybackEvent(msg))
          console.error('[relay] error:', msg.message)
          break
        }

        default:
          break
      }
    }

    ws.onclose = () => {
      cleanup()
      setState('disconnected')

      if (reconnectAttemptsRef.current < maxReconnect) {
        const delay = Math.min(2000 * 2 ** reconnectAttemptsRef.current, 60000)
        reconnectAttemptsRef.current++
        reconnectTimerRef.current = setTimeout(() => {
          connect()
        }, delay)
      } else {
        console.warn('[relay] max reconnect attempts reached, cooling down for 5 minutes...')
        reconnectTimerRef.current = setTimeout(() => {
          reconnectAttemptsRef.current = 0
          connect()
        }, 5 * 60_000)
      }
    }

    ws.onerror = () => {
      setState('error')
    }
  }, [token, cleanup, signOut])

  useEffect(() => {
    if (!token) return
    connect()
    return () => {
      cleanup()
      if (wsRef.current) {
        wsRef.current.close(1000, 'Provider unmount')
        wsRef.current = null
      }
    }
  }, [token, connect, cleanup])

  const joinSession = useCallback(
    (sessionId: string, handlers: React.RefObject<SessionHandlers>) => {
      sessionHandlersRef.current.set(sessionId, { handlers })
      send({ type: 'session_join', sessionId })
    },
    [send],
  )

  const leaveSession = useCallback(
    (sessionId: string) => {
      sessionHandlersRef.current.delete(sessionId)
      send({ type: 'session_leave', sessionId })
    },
    [send],
  )

  return createElement(
    RelayContext.Provider,
    { value: { state, avatarEvent, publishAvatarEvent, send, joinSession, leaveSession } },
    children,
  )
}

// ── Hooks ────────────────────────────────────────────────────────

export function useRelay(): RelayContextValue {
  const ctx = useContext(RelayContext)
  if (!ctx) {
    throw new Error('useRelay must be used within a RelayProvider')
  }
  return ctx
}

/**
 * Subscribe to a specific chat session via the global WebSocket.
 * Sends session_join on mount (+ after reconnect) and session_leave on unmount.
 */
export function useRelaySession(
  sessionId: string,
  handlers: SessionHandlers,
) {
  const { state, send, joinSession, leaveSession } = useRelay()

  // Keep handlers in a ref to avoid effect re-runs
  const handlersRef = useRef<SessionHandlers>(handlers)
  handlersRef.current = handlers

  useEffect(() => {
    if (!sessionId) return
    joinSession(sessionId, handlersRef)
    return () => {
      leaveSession(sessionId)
    }
  }, [sessionId, joinSession, leaveSession])

  const requestHistory = useCallback(
    (before?: number, limit?: number) => {
      if (!sessionId) return
      send({
        type: 'history_request',
        sessionId,
        before,
        limit,
      })
    },
    [send, sessionId],
  )

  const clearHistory = useCallback(() => {
    if (!sessionId) return
    send({ type: 'history_clear', sessionId })
  }, [send, sessionId])

  return { state, send, requestHistory, clearHistory }
}
