import Foundation

final class WebSocketChatClient: NSObject {
    private var session: URLSession?
    private var task: URLSessionWebSocketTask?
    private var heartbeatTimer: Timer?
    private var joinTimeoutWorkItem: DispatchWorkItem?
    private var roomId: Int = 0
    private var isClosing = false
    private let joinTimeout: TimeInterval = 15

    private(set) var state: UserCreatorChatSocketState = .disconnected {
        didSet {
            onStateChange?(state)
        }
    }

    var onStateChange: ((UserCreatorChatSocketState) -> Void)?
    var onEvent: ((UserCreatorChatSocketEvent) -> Void)?

    func connect(roomId: Int) {
        performOnMain {
            connectOnMain(roomId: roomId)
        }
    }

    func sendText(_ text: String) -> String? {
        return performOnMain {
            sendTextOnMain(text)
        }
    }

    func ping() {
        performOnMain {
            pingOnMain()
        }
    }

    func close() {
        performOnMain {
            closeOnMain()
        }
    }

    func closeWithoutLeave(notify: Bool = true) {
        performOnMain {
            closeWithoutLeaveOnMain(notify: notify)
        }
    }

    private func connectOnMain(roomId: Int) {
        guard roomId > 0 else { return }
        closeWithoutLeaveOnMain(notify: false)

        self.roomId = roomId
        isClosing = false
        state = .connecting

        guard let url = makeWebSocketURL() else {
            state = .disconnected
            onEvent?(.closed(nil))
            return
        }

        var request = URLRequest(url: url)
        request.setValue("Bearer \(UserDefaults.string(forKey: UserDefaultsKey.token))", forHTTPHeaderField: "Authorization")

        let delegateQueue = OperationQueue()
        delegateQueue.maxConcurrentOperationCount = 1
        let session = URLSession(configuration: .default, delegate: self, delegateQueue: delegateQueue)
        let task = session.webSocketTask(with: request)
        self.session = session
        self.task = task
        task.resume()
        receiveLoop(for: task)
    }

    private func sendTextOnMain(_ text: String) -> String? {
        guard state == .joined, roomId > 0 else { return nil }
        let requestId = UUID().uuidString
        send(type: .sendText, requestId: requestId, payload: ["textMessage": text])
        return requestId
    }

    private func pingOnMain() {
        guard state == .joined, roomId > 0 else { return }
        send(type: .ping, requestId: UUID().uuidString, payload: [:])
    }

    private func closeOnMain() {
        isClosing = true
        stopHeartbeat()
        cancelJoinTimeout()

        guard let task else {
            state = .disconnected
            return
        }

        if state == .joined, roomId > 0 {
            send(type: .leaveRoom, requestId: UUID().uuidString, payload: [:]) { [weak self] _ in
                self?.cancelTask(task)
            }
        } else {
            cancelTask(task)
        }
    }

    private func closeWithoutLeaveOnMain(notify: Bool) {
        isClosing = true
        stopHeartbeat()
        cancelJoinTimeout()
        if let task {
            cancelTask(task)
        } else {
            state = .disconnected
            if notify {
                onEvent?(.closed(nil))
            }
        }
    }

    private func joinRoomIfPossible() {
        guard state == .socketOpen, roomId > 0 else { return }
        state = .joining
        startJoinTimeout()
        send(type: .joinRoom, requestId: UUID().uuidString, payload: [:]) { [weak self] error in
            guard let error else { return }
            self?.handleJoinFailure(message: error.localizedDescription)
        }
    }

    private func send(
        type: UserCreatorChatSocketMessageType,
        requestId: String,
        payload: [String: Any],
        completion: ((Error?) -> Void)? = nil
    ) {
        guard roomId > 0 else {
            completion?(WebSocketChatClientSendError.invalidState)
            return
        }
        guard type == .joinRoom || state == .joined else {
            completion?(WebSocketChatClientSendError.invalidState)
            return
        }

        let dictionary: [String: Any] = [
            "type": type.rawValue,
            "requestId": requestId,
            "roomId": roomId,
            "payload": payload
        ]

        guard let data = try? JSONSerialization.data(withJSONObject: dictionary),
              let json = String(data: data, encoding: .utf8) else {
            completion?(WebSocketChatClientSendError.encodingFailed)
            return
        }

        DEBUG_LOG(json)
        guard let task else {
            completion?(WebSocketChatClientSendError.taskUnavailable)
            return
        }

        task.send(.string(json)) { error in
            DispatchQueue.main.async {
                if let error {
                    ERROR_LOG(error.localizedDescription)
                }
                completion?(error)
            }
        }
    }

    private func receiveLoop(for task: URLSessionWebSocketTask) {
        task.receive { [weak self] result in
            DispatchQueue.main.async {
                guard let self, self.isCurrentTask(task) else { return }

                switch result {
                case .success(let message):
                    self.handle(message)
                    self.receiveLoop(for: task)
                case .failure(let error):
                    self.finishCurrentTask(task, error: error)
                }
            }
        }
    }

    private func handle(_ message: URLSessionWebSocketTask.Message) {
        let rawText: String?
        switch message {
        case .string(let text):
            rawText = text
        case .data(let data):
            rawText = String(data: data, encoding: .utf8)
        @unknown default:
            rawText = nil
        }

        guard let rawText else { return }
        DEBUG_LOG(rawText)

        guard let data = rawText.data(using: .utf8),
              let envelope = try? JSONDecoder().decode(UserCreatorChatIncomingEnvelope.self, from: data) else {
            return
        }

        switch envelope.type {
        case UserCreatorChatSocketMessageType.joined.rawValue:
            guard state == .joining else { return }
            cancelJoinTimeout()
            state = .joined
            startHeartbeat()
            onEvent?(.joined)
        case UserCreatorChatSocketMessageType.sendAck.rawValue:
            if let requestId = envelope.requestId {
                onEvent?(.sendAck(requestId: requestId, message: envelope.payload?.message))
            }
        case UserCreatorChatSocketMessageType.message.rawValue:
            if let message = envelope.payload?.message {
                onEvent?(.message(message))
            }
        case UserCreatorChatSocketMessageType.pong.rawValue:
            DEBUG_LOG("PONG")
            onEvent?(.pong)
        case UserCreatorChatSocketMessageType.error.rawValue:
            if state != .joined {
                cancelJoinTimeout()
            }
            let messageKey = envelope.payload?.messageKey
            if let messageKey {
                ERROR_LOG(messageKey)
            }
            onEvent?(.error(messageKey: messageKey))
        default:
            break
        }
    }

    private func startHeartbeat() {
        DispatchQueue.main.async { [weak self] in
            guard let self, self.state == .joined else { return }
            self.stopHeartbeat()
            self.heartbeatTimer = Timer.scheduledTimer(withTimeInterval: 30, repeats: true) { [weak self] _ in
                self?.ping()
            }
        }
    }

    private func startJoinTimeout() {
        cancelJoinTimeout()

        let workItem = DispatchWorkItem { [weak self] in
            guard let self, self.state == .joining else { return }
            self.joinTimeoutWorkItem = nil
            self.handleJoinFailure(message: "JOINED timeout")
        }
        joinTimeoutWorkItem = workItem
        DispatchQueue.main.asyncAfter(deadline: .now() + joinTimeout, execute: workItem)
    }

    private func cancelJoinTimeout() {
        joinTimeoutWorkItem?.cancel()
        joinTimeoutWorkItem = nil
    }

    private func handleJoinFailure(message: String?) {
        guard state == .joining else { return }
        if let message {
            ERROR_LOG(message)
        }
        closeWithoutLeaveOnMain(notify: false)
        onEvent?(.joinFailed(message: message))
    }

    private func stopHeartbeat() {
        let invalidate = { [weak self] in
            self?.heartbeatTimer?.invalidate()
            self?.heartbeatTimer = nil
        }

        if Thread.isMainThread {
            invalidate()
        } else {
            DispatchQueue.main.async(execute: invalidate)
        }
    }

    private func cancelTask(_ task: URLSessionWebSocketTask) {
        cancelJoinTimeout()
        task.cancel(with: .normalClosure, reason: nil)
        if isCurrentTask(task) {
            self.task = nil
            session?.invalidateAndCancel()
            session = nil
            state = .disconnected
        }
    }

    private func finishCurrentTask(_ task: URLSessionWebSocketTask, error: Error?) {
        guard isCurrentTask(task) else { return }
        let shouldNotify = isClosing == false

        stopHeartbeat()
        cancelJoinTimeout()
        self.task = nil
        session?.invalidateAndCancel()
        session = nil
        state = .disconnected

        if shouldNotify {
            if let error {
                ERROR_LOG(error.localizedDescription)
            }
            onEvent?(.closed(error))
        }
    }

    private func isCurrentTask(_ task: URLSessionWebSocketTask) -> Bool {
        self.task === task
    }

    private func makeWebSocketURL() -> URL? {
        guard var components = URLComponents(string: BASE_URL), let scheme = components.scheme else {
            return nil
        }

        components.scheme = scheme == "http" ? "ws" : "wss"
        components.path = "/ws/v2/user-creator-chat"
        components.query = nil
        return components.url
    }

    private func performOnMain<T>(_ work: () -> T) -> T {
        if Thread.isMainThread {
            return work()
        }
        return DispatchQueue.main.sync(execute: work)
    }
}

private enum WebSocketChatClientSendError: LocalizedError {
    case invalidState
    case encodingFailed
    case taskUnavailable

    var errorDescription: String? {
        switch self {
        case .invalidState:
            return "WebSocket is not ready to send"
        case .encodingFailed:
            return "Failed to encode WebSocket message"
        case .taskUnavailable:
            return "WebSocket task is unavailable"
        }
    }
}

extension WebSocketChatClient: URLSessionWebSocketDelegate {
    func urlSession(
        _ session: URLSession,
        webSocketTask: URLSessionWebSocketTask,
        didOpenWithProtocol protocol: String?
    ) {
        DispatchQueue.main.async { [weak self] in
            guard let self, self.isCurrentTask(webSocketTask) else { return }
            self.state = .socketOpen
            self.joinRoomIfPossible()
        }
    }

    func urlSession(
        _ session: URLSession,
        webSocketTask: URLSessionWebSocketTask,
        didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
        reason: Data?
    ) {
        DispatchQueue.main.async { [weak self] in
            self?.finishCurrentTask(webSocketTask, error: nil)
        }
    }
}