import Foundation
import Network

enum NetworkFraming {
    static let headerSize = 4  // UInt32 for length
    static let maxMessageSize: UInt32 = 10_485_760  // 10 MB

    /// Frame data with 4-byte length prefix
    static func frameData(_ data: Data) -> Data {
        var length = UInt32(data.count).bigEndian
        var framedData = Data(bytes: &length, count: headerSize)
        framedData.append(data)
        return framedData
    }

    /// Send framed data over connection
    static func send(_ data: Data, over connection: NWConnection, completion: @escaping (Error?) -> Void) {
        let framedData = frameData(data)
        connection.send(content: framedData, completion: .contentProcessed { error in
            completion(error)
        })
    }

    /// Receive framed data from connection
    static func receive(from connection: NWConnection, completion: @escaping (Result<Data, Error>) -> Void) {
        // Read 4-byte header
        connection.receive(minimumIncompleteLength: headerSize, maximumLength: headerSize) { headerData, _, isComplete, error in
            if let error = error {
                completion(.failure(error))
                return
            }

            guard let headerData = headerData, headerData.count == headerSize else {
                completion(.failure(FramingError.incompleteHeader))
                return
            }

            let length = headerData.withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }

            guard length > 0, length <= maxMessageSize else {
                completion(.failure(FramingError.invalidMessageSize(length)))
                return
            }

            // Read payload
            connection.receive(minimumIncompleteLength: Int(length), maximumLength: Int(length)) { payloadData, _, _, error in
                if let error = error {
                    completion(.failure(error))
                    return
                }

                guard let payloadData = payloadData else {
                    completion(.failure(FramingError.incompletePayload))
                    return
                }

                completion(.success(payloadData))
            }
        }
    }

    /// Continuous receive loop
    static func startReceiving(from connection: NWConnection, onData: @escaping (Data) -> Void, onError: @escaping (Error) -> Void) {
        receive(from: connection) { result in
            switch result {
            case .success(let data):
                onData(data)
                startReceiving(from: connection, onData: onData, onError: onError)
            case .failure(let error):
                onError(error)
            }
        }
    }
}

enum FramingError: LocalizedError {
    case incompleteHeader
    case incompletePayload
    case invalidMessageSize(UInt32)

    var errorDescription: String? {
        switch self {
        case .incompleteHeader: return "Incomplete header"
        case .incompletePayload: return "Incomplete payload"
        case .invalidMessageSize(let size): return "Invalid size: \(size)"
        }
    }
}
