import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
import { renderHook, act, waitFor } from '@testing-library/react'
import { useWebSocket } from '../useWebSocket'

// Mock WebSocket class
class MockWebSocket {
    static CONNECTING = 0
    static OPEN = 1
    static CLOSING = 2
    static CLOSED = 3

    readyState = MockWebSocket.CONNECTING
    url: string
    onopen: (() => void) | null = null
    onmessage: ((event: { data: string }) => void) | null = null
    onerror: ((error: Event) => void) | null = null
    onclose: (() => void) | null = null

    constructor(url: string) {
        this.url = url
        mockWebSocketInstances.push(this)
    }

    send = vi.fn()
    close = vi.fn(() => {
        this.readyState = MockWebSocket.CLOSED
        this.onclose?.()
    })

    // Helper to simulate connection
    simulateOpen() {
        this.readyState = MockWebSocket.OPEN
        this.onopen?.()
    }

    // Helper to simulate message
    simulateMessage(data: object) {
        this.onmessage?.({ data: JSON.stringify(data) })
    }

    // Helper to simulate close
    simulateClose() {
        this.readyState = MockWebSocket.CLOSED
        this.onclose?.()
    }

    // Helper to simulate error
    simulateError() {
        this.onerror?.(new Event('error'))
    }
}

let mockWebSocketInstances: MockWebSocket[] = []

describe('useWebSocket', () => {
    beforeEach(() => {
        vi.useFakeTimers()
        mockWebSocketInstances = []
            ; (globalThis as any).WebSocket = MockWebSocket
    })

    afterEach(() => {
        vi.useRealTimers()
        vi.clearAllMocks()
    })

    it('initializes as disconnected', () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        expect(result.current.isConnected).toBe(false)
    })

    it('connects and sets isConnected to true', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        // Simulate WebSocket open
        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        await waitFor(() => {
            expect(result.current.isConnected).toBe(true)
        })
    })

    it('receives and stores messages', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        const testMessage = { type: 'order_update', order_id: '123' }
        act(() => {
            mockWebSocketInstances[0].simulateMessage(testMessage)
        })

        await waitFor(() => {
            expect(result.current.lastMessage).toEqual(expect.objectContaining(testMessage))
        })
    })

    it('adds messages to message history', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        const message1 = { type: 'order_update' as const, order_id: '1' }
        const message2 = { type: 'position_update' as const, position_id: '2' }

        act(() => {
            mockWebSocketInstances[0].simulateMessage(message1)
            mockWebSocketInstances[0].simulateMessage(message2)
        })

        await waitFor(() => {
            expect(result.current.messageHistory).toHaveLength(2)
        })
    })

    it('sends messages when connected', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        const testMessage = { action: 'subscribe', channel: 'prices' }
        act(() => {
            result.current.sendMessage(testMessage)
        })

        expect(mockWebSocketInstances[0].send).toHaveBeenCalledWith(JSON.stringify(testMessage))
    })

    it('allows subscribing to specific message types', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))
        const handler = vi.fn()

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
            result.current.subscribe('order_update', handler)
        })

        const testMessage = { type: 'order_update', order_id: '123' }
        act(() => {
            mockWebSocketInstances[0].simulateMessage(testMessage)
        })

        await waitFor(() => {
            expect(handler).toHaveBeenCalledWith(expect.objectContaining(testMessage))
        })
    })

    it('cleans up subscription on unsubscribe', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))
        const handler = vi.fn()

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        let unsubscribe: () => void
        act(() => {
            unsubscribe = result.current.subscribe('order_update', handler)
        })

        act(() => {
            unsubscribe()
        })

        const testMessage = { type: 'order_update', order_id: '123' }
        act(() => {
            mockWebSocketInstances[0].simulateMessage(testMessage)
        })

        // Handler should not be called after unsubscribe
        expect(handler).not.toHaveBeenCalled()
    })

    it('sets isConnected to false on disconnect', async () => {
        const { result } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        await waitFor(() => {
            expect(result.current.isConnected).toBe(true)
        })

        act(() => {
            mockWebSocketInstances[0].simulateClose()
        })

        await waitFor(() => {
            expect(result.current.isConnected).toBe(false)
        })
    })

    it('closes WebSocket on unmount', () => {
        const { unmount } = renderHook(() => useWebSocket('ws://localhost:8000/ws'))

        act(() => {
            mockWebSocketInstances[0].simulateOpen()
        })

        unmount()

        expect(mockWebSocketInstances[0].close).toHaveBeenCalled()
    })
})
