import Foundation
import Network
import os

/// Wraps an `NWConnection` for sending and receiving newline-delimited JSON messages.
@MainActor
final class NDJSONTransport {
    private let connection: NWConnection
    private let lineBuffer: NDJSONLineBuffer
    private let logger: os.Logger
    private let encoder: JSONEncoder = {
        let e = JSONEncoder()
        e.outputFormatting = [.sortedKeys]
        return e
    }()

    /// Called for each complete line received from the connection.
    var onLine: ((String) -> Void)?

    /// Called when the connection is closed or encounters an error.
    var onClose: (() -> Void)?

    init(connection: NWConnection, logger: os.Logger) {
        self.connection = connection
        self.logger = logger

        // Capture self weakly in the line buffer callback to avoid retain cycles.
        // The buffer is created before self is fully initialized, so we set the
        // actual forwarding closure after init completes via a two-phase approach.
        var forwardLine: ((String) -> Void)?
        self.lineBuffer = NDJSONLineBuffer { line in
            forwardLine?(line)
        }
        forwardLine = { [weak self] line in
            self?.onLine?(line)
        }
    }

    /// JSON-encode the message, append a newline, and send it over the connection.
    func send<T: Encodable>(_ message: T) {
        do {
            var data = try encoder.encode(message)
            data.append(contentsOf: [UInt8(ascii: "\n")])
            connection.send(content: data, completion: .contentProcessed { [weak self] error in
                if let error {
                    self?.logger.error("Send failed: \(error.localizedDescription)")
                }
            })
        } catch {
            logger.error("Encode failed: \(error.localizedDescription)")
        }
    }

    /// Begin the receive loop. Data is fed into the line buffer, which emits
    /// complete lines via the `onLine` callback.
    func startReceiving() {
        receiveNext()
    }

    /// Cancel the underlying connection.
    func close() {
        connection.cancel()
    }

    // MARK: - Private

    private func receiveNext() {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 65_536) { [weak self] content, _, isComplete, error in
            Task { @MainActor [weak self] in
                guard let self else { return }

                if let data = content, !data.isEmpty {
                    self.lineBuffer.feed(data)
                }

                if isComplete {
                    self.logger.info("Connection closed by peer")
                    self.onClose?()
                } else if let error {
                    self.logger.error("Receive error: \(error.localizedDescription)")
                    self.onClose?()
                } else {
                    self.receiveNext()
                }
            }
        }
    }
}
