/macos/RoonKit/Sources/RoonKit/Connection/Connection.swift at d8414c2e0ec68d181186716c9bf38201ee9b9ef7

// Copyright 2025 Shota FUJI
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

import Network
import OSLog
import Observation

public actor Connection: Connectable, Communicatable {
	private let logger = Logger(subsystem: subsystem, category: "Connection")

	public let serverID: String
	public let host: String
	public let port: UInt16
	public let ext: RegistryService.Extension

	public private(set) var token: String? = nil

	private var conn: NWConnection
	private var writeContext: NWConnection.ContentContext
	private let queue = DispatchQueue(label: "RoonKit.Connection.network")

	private var requestID: RequestId = 0

	private var messageContinuationID: UInt = 0
	private var messageContinuations: [UInt: AsyncStream<Moo>.Continuation] = [:]

	private var readLoop: Task<Void, any Error>? = nil

	public enum InitError: Error {
		case invalidHost
	}

	public init(
		id: String,
		host: String,
		port: UInt16,
		ext: RegistryService.Extension
	)
		async throws
	{
		self.serverID = id
		self.host = host
		self.port = port
		self.ext = ext
		self.token = ext.token

		let metadata = NWProtocolWebSocket.Metadata(opcode: .binary)
		self.writeContext = NWConnection.ContentContext(
			identifier: "writeBinaryContext",
			metadata: [metadata]
		)

		var url = URLComponents()
		url.scheme = "ws"
		url.host = host
		url.port = Int(port)
		url.path = "/api"

		guard let url = url.url else {
			throw InitError.invalidHost
		}

		let options = NWProtocolWebSocket.Options.init()
		options.autoReplyPing = true

		let parameters = NWParameters.tcp
		parameters.defaultProtocolStack.applicationProtocols.insert(options, at: 0)

		self.conn = NWConnection.init(to: .url(url), using: parameters)

		try await self.connect()
	}

	deinit {
		logger.info("Disconnecting... (automatic)")

		if let readLoop = readLoop {
			readLoop.cancel()
			self.readLoop = nil
		}

		self.messageContinuations = [:]

		if conn.state != .cancelled {
			conn.cancel()
		}
	}

	public enum ConnectError: Error {
		case alreadyStarted
		case unexpectedInfoReceived
		case serverIDMismatch
	}

	public func connect() async throws {
		switch conn.state {
		case .failed(let error):
			throw error
		case .cancelled, .setup:
			break
		default:
			throw ConnectError.alreadyStarted
		}

		logger.debug("Starting WebSocket connection")

		conn.start(queue: queue)
		try await conn.ready()
		logger.debug("WebSocket connection is ready")

		self.readLoop = startReadLoop()

		conn.stateUpdateHandler = { [weak self] (state: NWConnection.State) in
			switch state {
			case .cancelled, .failed(_), .waiting(_):
				Task {
					self?.logger.info("Network connection closed")
					await self?.disconnect()
				}
			case .preparing:
				self?.logger.debug("Network is in preparing state")
			case .setup:
				self?.logger.debug("Network has reset to setup state")
			default:
				break
			}
		}

		do {
			logger.debug("Querying server info")
			let res = try await request(Moo(info: .init()), timeout: .seconds(3))
			let info = try RegistryService.InfoResponse(res)
			guard info.coreId == serverID else {
				throw ConnectError.serverIDMismatch
			}
		}

		do {
			logger.debug("Registering extension")
			// This stays pending until user accepts the extension on Roon settings page.
			// Because of this, setting timeout does not make sense here.
			let res = try await request(Moo(register: ext))
			let result = try RegistryService.RegisterResponse(res)
			token = result.token
		}
	}

	private func clearMessageContinuations() {
		for (_, continuation) in messageContinuations {
			continuation.finish()
		}
		messageContinuations.removeAll()
	}

	private func startReadLoop() -> Task<Void, any Error> {
		// This loop task is a field of the actor containing it--every reference to
		// the actor (self) should be weak otherwise they form circular reference
		// and the actor will not be released via ARC.
		Task { [weak self] in
			while true {
				try Task.checkCancellation()

				do {
					guard let self = self else {
						return
					}

					let (data, context, _) = try await conn.receiveMessage()

					guard let data = data else {
						continue
					}

					guard
						let metadata = context?.protocolMetadata.first
							as? NWProtocolWebSocket.Metadata
					else {
						logger.warning(
							"Received non-WebSocket message on WebSocket connection"
						)
						continue
					}

					switch metadata.opcode {
					case .binary, .text:
						guard let msg = Moo.init(String(data: data, encoding: .utf8)!)
						else {
							logger.warning("Received invalid MOO message")
							continue
						}

						if msg.service == PingService.Ping.service {
							Task { [weak self] in
								var res = Moo(verb: "COMPLETE", service: "Success")
								res.requestId = msg.requestId

								guard let self = self else {
									return
								}

								try await send(res)
							}
							continue
						}

						for (_, continuation) in await messageContinuations {
							continuation.yield(msg)
						}
					case .close:
						logger.info("Connection closed, releasing resources")
						return
					default:
						break
					}
				} catch {
					if let logger = self?.logger {
						logger.warning("Receive failed: \(error)")
					}
					return
				}
			}
		}
	}

	public func disconnect() {
		logger.info("Disconnecting... (manual)")

		if let readLoop = readLoop {
			readLoop.cancel()
			self.readLoop = nil
		}

		self.messageContinuations = [:]

		if conn.state != .cancelled {
			conn.cancel()
		}
	}

	public var messages: AsyncStream<Moo> {
		let id = messageContinuationID
		messageContinuationID += 1

		return AsyncStream { cont in
			messageContinuations[id] = cont
			cont.onTermination = { @Sendable _ in
				Task.detached {
					await self.deleteMessageContinuation(id)
				}
			}
		}
	}

	private func request(
		msg: consuming Moo,
		timeout: ContinuousClock.Instant.Duration? = nil
	) async throws -> Moo {
		let requestID: RequestId
		if let givenID = msg.requestId {
			requestID = givenID
		} else {
			requestID = self.requestID
			msg.requestId = requestID
			self.requestID += 1
		}

		let timeoutTask: Task<Void, any Error>?
		if let timeout = timeout {
			timeoutTask = Task {
				try await Task.sleep(for: timeout)
				throw MessageReadError.timeout
			}
		} else {
			timeoutTask = nil
		}

		let read = Task {
			for await msg in self.messages {
				guard msg.requestId == .some(requestID) else {
					continue
				}

				timeoutTask?.cancel()
				return msg
			}

			throw MessageReadError.connectionClosed
		}

		try await self.send(msg)

		return try await read.value
	}

	public func request(_ msg: consuming Moo) async throws -> Moo {
		return try await request(msg: msg)
	}

	public func request(
		_ msg: consuming Moo,
		timeout: ContinuousClock.Instant.Duration
	) async throws -> Moo {
		return try await request(msg: msg, timeout: timeout)
	}

	public func send(_ msg: consuming Moo) async throws {
		let req = String(msg).data(using: .utf8)!

		try await withCheckedThrowingContinuation {
			(cont: CheckedContinuation<Void, any Error>) in
			if conn.state != .ready {
				cont.resume(throwing: MessageWriteError.connectionNotReady)
				return
			}

			conn.send(
				content: req,
				contentContext: writeContext,
				completion: .contentProcessed({ error in
					if let error = error {
						cont.resume(throwing: error)
						return
					}

					cont.resume()
				})
			)
		}
	}

	private func deleteMessageContinuation(_ id: UInt) {
		self.messageContinuations[id] = nil
	}
}

extension NWConnection {
	fileprivate enum ConnectionNotReadyError: Error {
		case failed(NWError)
		case cancelled
	}

	fileprivate func ready() async throws {
		try await withCheckedThrowingContinuation {
			(cont: CheckedContinuation<Void, any Error>) in
			self.stateUpdateHandler = { state in
				switch state {
				case .failed(let error), .waiting(let error):
					cont.resume(throwing: ConnectionNotReadyError.failed(error))
					self.stateUpdateHandler = nil
				case .cancelled:
					cont.resume(throwing: ConnectionNotReadyError.cancelled)
					self.stateUpdateHandler = nil
				case .ready:
					cont.resume()
					self.stateUpdateHandler = nil
				default:
					break
				}
			}
		}
	}

	fileprivate func receiveMessage() async throws -> (
		Data?, ContentContext?, Bool
	) {
		return try await withCheckedThrowingContinuation { cont in
			self.receiveMessage { (data, context, received, error) in
				if let error = error {
					cont.resume(throwing: error)
					return
				}

				cont.resume(returning: (data, context, received))
			}
		}
	}

	fileprivate func send(
		_ data: Data,
		contentContext: ContentContext = ContentContext.defaultMessage
	) async throws {
		try await withCheckedThrowingContinuation {
			(cont: CheckedContinuation<Void, any Error>) in
			self.send(
				content: data,
				contentContext: contentContext,
				completion: .contentProcessed({ error in
					if let error = error {
						cont.resume(throwing: error)
						return
					}

					cont.resume()
				})
			)
		}
	}
}