import Foundation
import Network

public actor DirectTransport: Transport {
    private let endpoint: TransportEndpoint
    private var connection: NWConnection?
    private var receiveContinuation: AsyncThrowingStream<Data, Error>.Continuation?

    public init(endpoint: TransportEndpoint) {
        self.endpoint = endpoint
    }

    public func connect() async throws {
        let host = NWEndpoint.Host(endpoint.host)
        let port = NWEndpoint.Port(integerLiteral: UInt16(endpoint.port))
        let connection = NWConnection(host: host, port: port, using: .tcp)
        self.connection = connection

        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            connection.stateUpdateHandler = { state in
                switch state {
                case .ready:
                    continuation.resume()
                case .failed(let error):
                    continuation.resume(throwing: error)
                case .cancelled:
                    continuation.resume(throwing: CancellationError())
                default:
                    break
                }
            }
            connection.start(queue: .global(qos: .userInitiated))
        }
        startReceiveLoop()
    }

    public func send(_ data: Data) async throws {
        guard let connection else { throw TransportError.notConnected }
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            connection.send(content: data, completion: .contentProcessed { error in
                if let error { continuation.resume(throwing: error) }
                else { continuation.resume() }
            })
        }
    }

    public nonisolated func receive() -> AsyncThrowingStream<Data, Error> {
        AsyncThrowingStream { continuation in
            Task { await self.installContinuation(continuation) }
        }
    }

    public func disconnect() async {
        connection?.cancel()
        connection = nil
        receiveContinuation?.finish()
        receiveContinuation = nil
    }

    private func installContinuation(_ continuation: AsyncThrowingStream<Data, Error>.Continuation) {
        self.receiveContinuation = continuation
    }

    private func startReceiveLoop() {
        guard let connection else { return }
        connection.receive(minimumIncompleteLength: 1, maximumLength: 65_536) { [weak self] data, _, isComplete, error in
            guard let self else { return }
            Task {
                await self.forward(data: data, isComplete: isComplete, error: error)
            }
        }
    }

    private func forward(data: Data?, isComplete: Bool, error: NWError?) {
        if let error {
            receiveContinuation?.finish(throwing: error)
            return
        }
        if let data, !data.isEmpty {
            receiveContinuation?.yield(data)
        }
        if isComplete {
            receiveContinuation?.finish()
            return
        }
        startReceiveLoop()
    }
}

public enum TransportError: Error, Sendable {
    case notConnected
}
