Kecalek_python/ios_client/EncryptedChat/Network/ConnectionManager.swift

import Foundation
import Network

/// TCP connection manager using Network.framework.
/// Handles connection lifecycle, TLS, buffered reading (newline-delimited), and writing.
actor ConnectionManager {

    enum ConnectionState: Equatable {
        case disconnected
        case connecting
        case connected
        case failed(String)
    }

    private var connection: NWConnection?
    private var receiveBuffer = Data()
    private(set) var state: ConnectionState = .disconnected
    private var stateCallback: ((ConnectionState) -> Void)?
    private var messageStream: AsyncStream<[String: Any]>.Continuation?

    /// Set a callback for connection state changes
    func onStateChange(_ callback: @escaping (ConnectionState) -> Void) {
        stateCallback = callback
    }

    // MARK: - Connect / Disconnect

    /// Connect to server
    func connect(host: String, port: UInt16, useTLS: Bool = false, tlsInsecure: Bool = false) async throws {
        guard state == .disconnected || state != .connected else {
            throw NetworkError.alreadyConnected
        }

        updateState(.connecting)

        let nwHost = NWEndpoint.Host(host)
        let nwPort = NWEndpoint.Port(rawValue: port)!

        let params: NWParameters
        if useTLS {
            let tlsOptions = NWProtocolTLS.Options()
            if tlsInsecure {
                // Skip certificate verification (dev only)
                sec_protocol_options_set_verify_block(
                    tlsOptions.securityProtocolOptions,
                    { _, _, completionHandler in completionHandler(true) },
                    .main
                )
            }
            params = NWParameters(tls: tlsOptions, tcp: .init())
        } else {
            params = .tcp
        }

        let conn = NWConnection(host: nwHost, port: nwPort, using: params)
        self.connection = conn
        self.receiveBuffer = Data()

        return try await withCheckedThrowingContinuation { continuation in
            conn.stateUpdateHandler = { [weak self] newState in
                Task { [weak self] in
                    guard let self = self else { return }
                    switch newState {
                    case .ready:
                        await self.updateState(.connected)
                        continuation.resume()
                    case .failed(let error):
                        await self.updateState(.failed(error.localizedDescription))
                        continuation.resume(throwing: NetworkError.connectionFailed(error.localizedDescription))
                    case .cancelled:
                        await self.updateState(.disconnected)
                    case .waiting(let error):
                        await self.updateState(.failed(error.localizedDescription))
                        continuation.resume(throwing: NetworkError.connectionFailed("Waiting: \(error.localizedDescription)"))
                    default:
                        break
                    }
                }
            }
            conn.start(queue: .global(qos: .userInitiated))
        }
    }

    /// Disconnect from server
    func disconnect() {
        connection?.cancel()
        connection = nil
        receiveBuffer = Data()
        updateState(.disconnected)
        messageStream?.finish()
        messageStream = nil
    }

    // MARK: - Send

    /// Send raw data over the connection
    func send(_ data: Data) async throws {
        guard let connection = connection, state == .connected else {
            throw NetworkError.notConnected
        }

        return try await withCheckedThrowingContinuation { continuation in
            connection.send(content: data, completion: .contentProcessed { error in
                if let error = error {
                    continuation.resume(throwing: NetworkError.connectionFailed(error.localizedDescription))
                } else {
                    continuation.resume()
                }
            })
        }
    }

    /// Send a protocol message (builds JSON + newline, sends)
    func sendMessage(type: String, requestId: String? = nil, params: [String: Any] = [:]) async throws {
        let data = try ProtocolHandler.buildRequest(type: type, requestId: requestId, params: params)
        try await send(data)
    }

    // MARK: - Receive

    /// Read one newline-delimited JSON message.
    /// Returns nil on EOF / connection close.
    func readMessage() async throws -> [String: Any]? {
        while true {
            // Check buffer for a complete line
            if let newlineIndex = receiveBuffer.firstIndex(of: 0x0A) {
                let lineData = receiveBuffer.prefix(through: newlineIndex)
                receiveBuffer.removeSubrange(...newlineIndex)

                // Check size
                if lineData.count > Constants.maxMessageBytes {
                    throw NetworkError.messageTooLarge
                }

                return try ProtocolHandler.parseMessage(Data(lineData))
            }

            // Buffer doesn't have a complete line — read more from the connection
            guard let connection = connection else {
                return nil
            }

            let chunk = try await receiveChunk(connection: connection)
            guard let chunk = chunk else {
                return nil  // EOF
            }

            receiveBuffer.append(chunk)

            // Safety: if buffer exceeds max without a newline, drop it
            if receiveBuffer.count > Constants.maxMessageBytes * 2 {
                receiveBuffer = Data()
                throw NetworkError.messageTooLarge
            }
        }
    }

    /// Read a chunk of data from the connection
    private func receiveChunk(connection: NWConnection) async throws -> Data? {
        return try await withCheckedThrowingContinuation { continuation in
            connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) { content, _, isComplete, error in
                if let error = error {
                    continuation.resume(throwing: NetworkError.connectionFailed(error.localizedDescription))
                    return
                }
                if let content = content, !content.isEmpty {
                    continuation.resume(returning: content)
                } else if isComplete {
                    continuation.resume(returning: nil)
                } else {
                    // No data and not complete — shouldn't happen but return nil
                    continuation.resume(returning: nil)
                }
            }
        }
    }

    // MARK: - State

    var isConnected: Bool {
        state == .connected
    }

    private func updateState(_ newState: ConnectionState) {
        state = newState
        stateCallback?(newState)
    }
}