import Foundation
import Combine
import UIKit

/// Transport layer for WebSocket communication.
/// Manages the raw connection, reconnection with exponential backoff, and message sending.
class WebSocketConnection: ObservableObject {
    @Published var isConnected: Bool = false

    private var webSocketTask: URLSessionWebSocketTask?
    private let urlSession: URLSession
    private var connectionAttemptDelay: TimeInterval = 0.1
    private let retryMaxAttempts = 1000
    private var retryCurrentAttempt = 0
    private var isReceiving = false

    /// Called when a text message is received from the server.
    var onMessage: ((String) -> Void)?

    /// Called when a connection error occurs (for analytics/error state).
    var onError: (() -> Void)?

    init(urlSession: URLSession) {
        self.urlSession = urlSession
        self.webSocketTask = self.urlSession.webSocketTask(with: URL(string: wsEnvironment.url)!)
        self.webSocketTask?.resume()
        #if DEBUG
        print("WebSocket initialized with \(wsEnvironment) environment: \(wsEnvironment.url)")
        #endif
    }

    func connect() {
        #if DEBUG
        print("connecting in \(connectionAttemptDelay) seconds to \(wsEnvironment.url)")
        #endif

        DispatchQueue.main.asyncAfter(deadline: .now() + connectionAttemptDelay) {
            self.webSocketTask = self.urlSession.webSocketTask(with: URL(string: wsEnvironment.url)!)
            self.webSocketTask?.resume()
            self.sendMessage(message: WebSocketCommands.initSession(
                userID: UIDevice.current.identifierForVendor?.uuidString ?? "Unknown"
            ))
            self.readMessage()
        }
    }

    func disconnect() {
        #if DEBUG
        print("disconnecting")
        #endif
        self.webSocketTask?.cancel(with: .goingAway, reason: nil)
        self.isReceiving = false
        self.isConnected = false
    }

    func sendMessage(message: [String: Any]) {
        guard let webSocketTask = self.webSocketTask else {
            #if DEBUG
            print("WebSocket task is nil")
            #endif
            return
        }

        do {
            #if DEBUG
            print(message)
            #endif
            let messageData = try JSONSerialization.data(withJSONObject: message, options: [])
            webSocketTask.send(.data(messageData)) { [weak self] error in
                guard let self = self else { return }
                if let error = error {
                    DispatchQueue.main.async {
                        self.isConnected = false
                        self.handleError()
                    }
                }
            }
        } catch {
            #if DEBUG
            print("Failed to encode message: \(error)")
            #endif
            DispatchQueue.main.async {
                self.isConnected = false
                self.handleError()
            }
        }
    }

    func readMessage() {
        guard !isReceiving, let webSocketTask = self.webSocketTask else { return }
        isReceiving = true

        webSocketTask.receive(completionHandler: { [weak self] result in
            guard let self = self else { return }

            DispatchQueue.main.async {
                self.isReceiving = false

                switch result {
                case .failure(let error):
                    self.isConnected = false

                    AnalyticsManager.shared.recordError(error, additionalInfo: [
                        "context": "websocket_receive",
                        "error_description": error.localizedDescription
                    ])
                    AnalyticsManager.shared.track(.websocketError(
                        errorType: "receive_failed",
                        message: error.localizedDescription
                    ))

                case .success(let message):
                    self.connectionEstablished()

                    switch message {
                    case .string(let text):
                        self.onMessage?(text)
                    case .data(let data):
                        #if DEBUG
                        print("Received data: \(data)")
                        #endif
                    @unknown default:
                        break
                    }

                    if self.isConnected {
                        self.readMessage()
                    }
                }
            }
        })
    }

    // MARK: - Private

    private func connectionEstablished() {
        isConnected = true
        retryCurrentAttempt = 0
        connectionAttemptDelay = 0.1
    }

    private func handleError() {
        isConnected = false

        SoundManager.shared.playSound(.error)

        let error = NSError(
            domain: "WebSocketService",
            code: -1,
            userInfo: [NSLocalizedDescriptionKey: "WebSocket connection lost"]
        )
        AnalyticsManager.shared.recordError(error, additionalInfo: [
            "context": "websocket_connection",
            "retry_attempt": retryCurrentAttempt,
            "max_retries": retryMaxAttempts,
            "connection_delay": connectionAttemptDelay
        ])
        AnalyticsManager.shared.track(.websocketDisconnected(reason: "connection_error"))

        onError?()

        guard retryCurrentAttempt < retryMaxAttempts else {
            #if DEBUG
            print("Could not reconnect after \(retryCurrentAttempt) attempts")
            #endif
            return
        }

        connectionAttemptDelay = min(connectionAttemptDelay * 1.1, 5)

        DispatchQueue.main.asyncAfter(deadline: .now() + connectionAttemptDelay) {
            #if DEBUG
            print("attempting reconnect # \(self.retryCurrentAttempt)")
            #endif
            self.reconnect()
        }
    }

    private func reconnect() {
        if webSocketTask?.state != .running {
            retryCurrentAttempt += 1
            connect()
        }
    }
}