import { ref } from 'vue'

export type SessionStatus = 'disconnected' | 'authenticating' | 'spinning_up' | 'initializing' | 'ready' | 'error'

export interface WebSocketMessage {
  type: string
  [key: string]: any
}

export type MessageHandler = (message: WebSocketMessage) => void

class WebSocketManager {
  private ws: WebSocket | null = null
  private handlers: Set<MessageHandler> = new Set()
  private reconnectTimeout: number | null = null
  public isConnected = ref(false)
  public isAuthenticated = ref(false)
  public sessionStatus = ref<SessionStatus>('disconnected')
  public statusMessage = ref('')
  private token: string | null = null
  private messageQueue: WebSocketMessage[] = []
  private reconnectAttempts = 0
  private maxReconnectAttempts = Infinity // Keep trying indefinitely
  private reconnectDelay = 1000 // Start with 1 second
  private maxReconnectDelay = 50000 // Max 50 seconds
  private consolePatchInstalled = false
  private forwardingLog = false // prevent re-entrancy

  /**
   * Connect to WebSocket with JWT token for authentication
   */
  async connect(token: string): Promise<void> {
    this.token = token

    // Close existing connection if any — null out handlers first so the async
    // onclose event from the old socket cannot reset sessionStatus after the
    // new socket has already reached 'ready'.
    if (this.ws) {
      this.ws.onopen = null
      this.ws.onmessage = null
      this.ws.onerror = null
      this.ws.onclose = null
      this.ws.close()
      this.ws = null
    }

    return new Promise((resolve, reject) => {
      // Build WebSocket URL based on environment
      // VITE_GATEWAY_URL can be:
      //   - Relative path: "/api" (for production with ingress)
      //   - Absolute URL: "http://localhost:3000" (for local dev)
      const envUrl = import.meta.env.VITE_GATEWAY_URL
      let baseUrl: string

      if (envUrl) {
        // If it starts with /, it's a relative path - use current origin
        if (envUrl.startsWith('/')) {
          baseUrl = window.location.origin + envUrl
        } else {
          baseUrl = envUrl
        }
      } else if (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1') {
        baseUrl = 'http://localhost:3000'
      } else {
        // Production: use current origin
        baseUrl = window.location.origin
      }

      // Gateway WebSocket endpoint is /ws/chat (baseUrl already includes /api if needed)
      const wsUrl = baseUrl.replace('http://', 'ws://').replace('https://', 'wss://') + '/ws/chat'

      // WebSocket API doesn't support custom headers, so we'll send the token in the URL
      // The gateway will need to handle token extraction from query params
      const wsUrlWithToken = `${wsUrl}?token=${encodeURIComponent(token)}`
      console.log('[WebSocket] Connecting to:', wsUrl)
      console.log('[WebSocket] Full URL (token hidden):', wsUrl + '?token=***')
      this.ws = new WebSocket(wsUrlWithToken)

      this.ws.onopen = () => {
        console.log('[WebSocket] Connected successfully')
        // Cancel any pending reconnect timer — we're already connected
        if (this.reconnectTimeout) {
          clearTimeout(this.reconnectTimeout)
          this.reconnectTimeout = null
        }
        this.isConnected.value = true
        this.isAuthenticated.value = false // Wait for 'connected' message from server
        this.reconnectAttempts = 0 // Reset reconnection counter
        resolve()
      }

      this.ws.onmessage = (event) => {
        console.log('[WebSocket] Raw message received:', event.data.substring(0, 200))
        try {
          const message = JSON.parse(event.data)
          console.log('[WebSocket] Parsed message type:', message.type)

          if (message.type === 'status') {
            this.sessionStatus.value = message.status as SessionStatus
            this.statusMessage.value = message.message ?? ''
          } else if (message.type === 'connected') {
            console.log('[WebSocket] Received connected message, marking as authenticated')
            this.isAuthenticated.value = true
            this.sessionStatus.value = 'ready'
            this.statusMessage.value = ''
            this.installConsoleForwarding(message.sessionId)
            // Flush any queued messages now that we're authenticated
            this.flushMessageQueue()
          } else if (message.type === 'error') {
            this.sessionStatus.value = 'error'
            this.statusMessage.value = message.message ?? ''
          }

          // Pass to all handlers
          this.handlers.forEach(handler => handler(message))
        } catch (err) {
          console.error('[WebSocket] Error parsing message:', err)
          console.error('[WebSocket] Raw data:', event.data.substring(0, 500))
        }
      }

      this.ws.onerror = (error) => {
        console.error('[WebSocket] Connection error:', error)
        this.isConnected.value = false
        reject(error)
      }

      this.ws.onclose = (event) => {
        this.isConnected.value = false
        this.isAuthenticated.value = false
        this.sessionStatus.value = 'disconnected'
        this.statusMessage.value = ''
        console.log('WebSocket disconnected:', event.code, event.reason)

        // Attempt to reconnect if we have a token and it wasn't an intentional close.
        // Check code instead of wasClean: code 1005 has wasClean=true but still needs retry.
        if (this.token && event.code !== 1000 && event.code !== 1001) {
          this.scheduleReconnect()
        }
      }

      // Connection timeout
      setTimeout(() => {
        if (!this.isConnected.value) {
          reject(new Error('WebSocket connection timeout'))
          this.ws?.close()
        }
      }, 10000)
    })
  }

  addHandler(handler: MessageHandler) {
    this.handlers.add(handler)
  }

  removeHandler(handler: MessageHandler) {
    this.handlers.delete(handler)
  }

  send(message: WebSocketMessage) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      console.log('[WebSocket] Sending message:', message.type)
      this.ws.send(JSON.stringify(message))
    } else {
      console.log('[WebSocket] Queuing message (not connected yet):', message.type, '- Queue size:', this.messageQueue.length + 1)
      this.messageQueue.push(message)
      // Only schedule reconnect if socket is CLOSED/undefined — not when it's still CONNECTING
      if (this.token && !this.reconnectTimeout && this.ws?.readyState !== WebSocket.CONNECTING) {
        this.scheduleReconnect()
      }
    }
  }

  private scheduleReconnect() {
    if (this.reconnectTimeout) {
      return // Reconnection already scheduled
    }

    // Exponential backoff with max delay
    const delay = Math.min(
      this.reconnectDelay * Math.pow(2, this.reconnectAttempts),
      this.maxReconnectDelay
    )

    console.log(`[WebSocket] Scheduling reconnect in ${delay}ms (attempt ${this.reconnectAttempts + 1})`)

    this.reconnectTimeout = window.setTimeout(() => {
      this.reconnectTimeout = null
      this.reconnectAttempts++

      if (this.token) {
        console.log('[WebSocket] Attempting to reconnect...')
        this.connect(this.token)
          .then(() => {
            console.log('[WebSocket] Reconnected successfully')
            this.reconnectAttempts = 0 // Reset counter on success
            this.flushMessageQueue()
          })
          .catch(err => {
            console.error('[WebSocket] Reconnection failed:', err)
            // scheduleReconnect will be called again via onclose
          })
      }
    }, delay)
  }

  private flushMessageQueue() {
    if (this.messageQueue.length > 0) {
      console.log(`[WebSocket] Flushing ${this.messageQueue.length} queued messages`)
      const queue = [...this.messageQueue]
      this.messageQueue = []
      queue.forEach(msg => this.send(msg))
    }
  }

  private installConsoleForwarding(sessionId: string) {
    if (this.consolePatchInstalled) return
    this.consolePatchInstalled = true

    const levels = ['log', 'info', 'warn', 'error', 'debug'] as const
    for (const level of levels) {
      const original = (console[level] as (...args: any[]) => void).bind(console)
      console[level] = (...args: any[]) => {
        original(...args)
        if (!this.forwardingLog && this.ws?.readyState === WebSocket.OPEN) {
          this.forwardingLog = true
          try {
            const message = args.map(a =>
              typeof a === 'string' ? a : (() => { try { return JSON.stringify(a) } catch { return String(a) } })()
            ).join(' ')
            this.ws!.send(JSON.stringify({ type: 'client_log', level, message }))
          } catch { /* ignore send errors */ }
          this.forwardingLog = false
        }
      }
    }
  }

  disconnect() {
    if (this.reconnectTimeout) {
      clearTimeout(this.reconnectTimeout)
      this.reconnectTimeout = null
    }
    this.token = null // Clear token to prevent auto-reconnect
    this.messageQueue = [] // Clear message queue
    this.reconnectAttempts = 0
    if (this.ws) {
      this.ws.close()
      this.ws = null
    }
    this.isAuthenticated.value = false
    this.sessionStatus.value = 'disconnected'
    this.statusMessage.value = ''
  }

  getWebSocket() {
    return this.ws
  }
}

export const wsManager = new WebSocketManager()
