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) 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 tlsOptions = NWProtocolTLS.Options()
        let params = NWParameters(tls: tlsOptions, tcp: .init())

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

        return try await withCheckedThrowingContinuation { continuation in
            // nonisolated flag — accessed only from the stateUpdateHandler serial queue
            // Use a class wrapper so the closure can mutate it
            final class ResumedFlag: @unchecked Sendable {
                var value = false
            }
            let resumed = ResumedFlag()

            conn.stateUpdateHandler = { [weak self] newState in
                Task { [weak self] in
                    guard let self = self else { return }
                    switch newState {
                    case .ready:
                        await self.updateState(.connected)
                        guard !resumed.value else { return }
                        resumed.value = true
                        continuation.resume()
                    case .failed(let error):
                        await self.updateState(.failed(error.localizedDescription))
                        guard !resumed.value else { return }
                        resumed.value = true
                        continuation.resume(throwing: NetworkError.connectionFailed(error.localizedDescription))
                    case .cancelled:
                        await self.updateState(.disconnected)
                        guard !resumed.value else { return }
                        resumed.value = true
                        continuation.resume(throwing: NetworkError.connectionFailed("Connection cancelled"))
                    case .waiting(let error):
                        await self.updateState(.failed(error.localizedDescription))
                        guard !resumed.value else { return }
                        resumed.value = true
                        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)
    }
}