import Foundation

actor WSClient {
    private let url: URL
    private let session: URLSession
    private let incomingStream: AsyncStream<WSMessage>
    private let incomingContinuation: AsyncStream<WSMessage>.Continuation

    private var webSocketTask: URLSessionWebSocketTask?
    private var receiveTask: Task<Void, Never>?
    private var pingTask: Task<Void, Never>?
    private var reconnectTask: Task<Void, Never>?
    private var reconnectID: UUID?
    private var authenticationTimedOutTask: URLSessionWebSocketTask?
    private var authToken: String?
    private var isAuthenticated = false
    private var isDisconnecting = false

    init(url: URL) {
        self.url = url
        self.session = URLSession(configuration: .default)

        let stream = AsyncStream<WSMessage>.makeStream()
        self.incomingStream = stream.stream
        self.incomingContinuation = stream.continuation
    }

    deinit {
        webSocketTask?.cancel(with: .goingAway, reason: nil)
        receiveTask?.cancel()
        pingTask?.cancel()
        reconnectTask?.cancel()
        incomingContinuation.finish()
    }

    var incoming: AsyncStream<WSMessage> {
        incomingStream
    }

    func connect(token: String) async throws {
        disconnect()
        isDisconnecting = false
        authToken = token

        try await establishConnection(token: token)
    }

    func disconnect() {
        isDisconnecting = true
        authToken = nil
        isAuthenticated = false
        reconnectID = nil
        authenticationTimedOutTask = nil

        receiveTask?.cancel()
        receiveTask = nil
        pingTask?.cancel()
        pingTask = nil
        reconnectTask?.cancel()
        reconnectTask = nil

        webSocketTask?.cancel(with: .goingAway, reason: nil)
        webSocketTask = nil
    }

    func send(_ message: WSMessage) async throws {
        guard isAuthenticated, let task = webSocketTask else {
            throw WSClientError.notConnected
        }

        let encodedMessage: String
        do {
            let data = try JSONEncoder().encode(message)
            guard let string = String(data: data, encoding: .utf8) else {
                throw WSClientError.encodingError
            }
            encodedMessage = string
        } catch let error as WSClientError {
            throw error
        } catch {
            throw WSClientError.encodingError
        }

        do {
            try await task.send(.string(encodedMessage))
        } catch {
            connectionDidFail(for: task)
            throw WSClientError.transportError(error)
        }
    }

    private func establishConnection(token: String) async throws {
        let task = session.webSocketTask(with: url)
        webSocketTask = task
        task.resume()

        do {
            try await sendAuthentication(token: token, using: task)
            let authResponse = try await receiveAuthenticationMessage(from: task)

            guard authResponse.type == "auth.ok" else {
                throw WSClientError.authenticationFailed(authResponse.reason)
            }

            guard webSocketTask === task, !isDisconnecting else {
                throw WSClientError.notConnected
            }

            isAuthenticated = true
            startReceiveLoop(for: task)
            startPingLoop(for: task)
        } catch let error as WSClientError {
            closeFailedConnection(task)
            throw error
        } catch {
            closeFailedConnection(task)
            throw WSClientError.transportError(error)
        }
    }

    private func sendAuthentication(token: String, using task: URLSessionWebSocketTask) async throws {
        let authMessage = WSMessage(type: "auth", token: token)
        let data = try JSONEncoder().encode(authMessage)

        guard let string = String(data: data, encoding: .utf8) else {
            throw WSClientError.encodingError
        }

        try await task.send(.string(string))
    }

    private func receiveAuthenticationMessage(
        from task: URLSessionWebSocketTask
    ) async throws -> WSMessage {
        let timeoutTask = Task { [weak self] in
            do {
                try await Task.sleep(nanoseconds: 10_000_000_000)
            } catch {
                return
            }

            await self?.timeoutAuthentication(for: task)
        }
        defer { timeoutTask.cancel() }

        do {
            let message = try await task.receive()
            return try decode(message)
        } catch {
            if authenticationTimedOutTask === task {
                authenticationTimedOutTask = nil
                throw WSClientError.authenticationTimedOut
            }
            throw error
        }
    }

    private func timeoutAuthentication(for task: URLSessionWebSocketTask) {
        guard webSocketTask === task, !isAuthenticated else {
            return
        }

        authenticationTimedOutTask = task
        task.cancel(with: .policyViolation, reason: nil)
    }

    private func startReceiveLoop(for task: URLSessionWebSocketTask) {
        receiveTask?.cancel()
        receiveTask = Task { [weak self] in
            await self?.receiveMessages(from: task)
        }
    }

    private func receiveMessages(from task: URLSessionWebSocketTask) async {
        do {
            while !Task.isCancelled, webSocketTask === task {
                let message = try await task.receive()
                incomingContinuation.yield(try decode(message))
            }
        } catch is CancellationError {
            return
        } catch {
            connectionDidFail(for: task)
        }
    }

    private func startPingLoop(for task: URLSessionWebSocketTask) {
        pingTask?.cancel()
        pingTask = Task { [weak self] in
            while !Task.isCancelled {
                do {
                    try await Task.sleep(nanoseconds: 30_000_000_000)
                } catch {
                    return
                }

                await self?.sendPing(using: task)
            }
        }
    }

    private func sendPing(using task: URLSessionWebSocketTask) async {
        guard isAuthenticated, webSocketTask === task else {
            return
        }

        do {
            try await task.send(.string("{\"type\":\"ping\"}"))
        } catch {
            connectionDidFail(for: task)
        }
    }

    private func connectionDidFail(for task: URLSessionWebSocketTask) {
        guard webSocketTask === task, !isDisconnecting else {
            return
        }

        isAuthenticated = false
        authenticationTimedOutTask = nil
        webSocketTask = nil
        receiveTask?.cancel()
        receiveTask = nil
        pingTask?.cancel()
        pingTask = nil
        task.cancel(with: .abnormalClosure, reason: nil)

        scheduleReconnect()
    }

    private func scheduleReconnect() {
        guard reconnectTask == nil, let token = authToken, !isDisconnecting else {
            return
        }

        let id = UUID()
        reconnectID = id
        reconnectTask = Task { [weak self] in
            await self?.reconnect(using: token, id: id)
        }
    }

    private func reconnect(using token: String, id: UUID) async {
        let delays: [UInt64] = [2, 4, 8, 16, 32]

        for delay in delays {
            do {
                try await Task.sleep(nanoseconds: delay * 1_000_000_000)
            } catch {
                return
            }

            guard reconnectID == id, !isDisconnecting else {
                return
            }

            do {
                try await establishConnection(token: token)
                if reconnectID == id {
                    reconnectID = nil
                    reconnectTask = nil
                }
                return
            } catch {
                continue
            }
        }

        if reconnectID == id {
            reconnectID = nil
            reconnectTask = nil
        }
    }

    private func closeFailedConnection(_ task: URLSessionWebSocketTask) {
        guard webSocketTask === task else {
            return
        }

        task.cancel(with: .policyViolation, reason: nil)
        webSocketTask = nil
        isAuthenticated = false
        authenticationTimedOutTask = nil
    }

    private func decode(_ message: URLSessionWebSocketTask.Message) throws -> WSMessage {
        let data: Data

        switch message {
        case .data(let messageData):
            data = messageData
        case .string(let string):
            data = Data(string.utf8)
        @unknown default:
            throw WSClientError.decodingError
        }

        do {
            return try JSONDecoder().decode(WSMessage.self, from: data)
        } catch {
            throw WSClientError.decodingError
        }
    }
}

struct WSMessage: Codable, Equatable {
    var type: String
    var id: String?
    var body: String?
    var to: String?
    var from: String?
    var fromName: String?
    var token: String?
    var peerID: String?
    var reason: String?
    var sentAt: String?          // ISO-8601 timestamp set by the sender; preserved through the relay
    var messageIDs: [String]?   // used by read receipts: { type: "read", messageIDs: [...] }
    var attachment: WSAttachment? // file metadata attached to a chat.message

    init(
        type: String,
        id: String? = nil,
        body: String? = nil,
        to: String? = nil,
        from: String? = nil,
        fromName: String? = nil,
        token: String? = nil,
        peerID: String? = nil,
        reason: String? = nil,
        sentAt: String? = nil,
        messageIDs: [String]? = nil,
        attachment: WSAttachment? = nil
    ) {
        self.type = type
        self.id = id
        self.body = body
        self.to = to
        self.from = from
        self.fromName = fromName
        self.token = token
        self.peerID = peerID
        self.reason = reason
        self.sentAt = sentAt
        self.messageIDs = messageIDs
        self.attachment = attachment
    }
}

struct WSAttachment: Codable, Equatable {
    let fileID: String
    let name: String
    let size: Int
}

enum WSClientError: Error {
    case notConnected
    case authenticationTimedOut
    case authenticationFailed(String?)
    case transportError(Error)
    case encodingError
    case decodingError
}
