Refactor HotlineFilePreviewClient to use NetSocketNew. Code is much cleaner.

2 files changed, 120 insertions, 145 deletions

diff --git a/Hotline/Hotline/HotlineTransferClient.swift b/Hotline/Hotline/HotlineTransferClient.swift
index 40dd598..86a1814 100644
--- a/Hotline/Hotline/HotlineTransferClient.swift
+++ b/Hotline/Hotline/HotlineTransferClient.swift
@@ -456,9 +456,9 @@ class HotlineFilePreviewClient: HotlineTransferClient {
   let serverPort: NWEndpoint.Port
   let referenceNumber: UInt32
   let referenceDataSize: UInt32
-  
+
   weak var delegate: HotlineFilePreviewClientDelegate? = nil
-  
+
   var status: HotlineTransferStatus = .unconnected {
     didSet {
       DispatchQueue.main.async {
@@ -466,156 +466,87 @@ class HotlineFilePreviewClient: HotlineTransferClient {
       }
     }
   }
-  
-  private var connection: NWConnection?
-  private var transferStage: HotlineFileTransferStage = .fileHeader
-  private var fileBytes = Data()
-  private var fileBytesTransferred: Int = 0
-  
+
+  private var downloadTask: Task<Void, Never>?
+
   init(address: String, port: UInt16, reference: UInt32, size: UInt32) {
     self.serverAddress = NWEndpoint.Host(address)
     self.serverPort = NWEndpoint.Port(rawValue: port + 1)!
     self.referenceNumber = reference
     self.referenceDataSize = size
   }
-  
+
   deinit {
-    self.invalidate()
+    downloadTask?.cancel()
   }
-  
+
   func start() {
-    guard self.status == .unconnected else {
+    guard status == .unconnected else {
       return
     }
-    
-    self.connect()
+
+    downloadTask = Task {
+      await self.download()
+    }
   }
-  
+
   func cancel() {
-    self.delegate = nil
-    
-    if self.status == .unconnected {
-      return
-    }
-    
-    self.invalidate()
-    
+    downloadTask?.cancel()
+    downloadTask = nil
+    delegate = nil
+
     print("HotlineFilePreviewClient: Cancelled preview transfer")
   }
-  
-  private func invalidate() {
-    if let c = self.connection {
-      c.stateUpdateHandler = nil
-      c.cancel()
-      
-      self.connection = nil
-    }
-    
-    self.fileBytes = Data()
-  }
-  
-  private func connect() {
-    self.connection = NWConnection(host: self.serverAddress, port: self.serverPort, using: .tcp)
-    self.connection?.stateUpdateHandler = { [weak self] (newState: NWConnection.State) in
-      switch newState {
-      case .ready:
-        self?.status = .connected
-        self?.sendMagic()
-      case .waiting(let err):
-        print("HotlineFilePreviewClient: Waiting", err)
-      case .cancelled:
-        print("HotlineFilePreviewClient: Cancelled")
-        self?.invalidate()
-      case .failed(let err):
-        print("HotlineFilePreviewClient: Connection error \(err)")
-        switch self?.status {
-        case .connecting:
-          print("HotlineFilePreviewClient: Failed to connect to file transfer server.")
-          self?.invalidate()
-          self?.status = .failed(.failedToConnect)
-        case .connected, .progress(_):
-          print("HotlineFilePreviewClient: Failed to finish transfer.")
-          self?.invalidate()
-          self?.status = .failed(.failedToDownload)
-        default:
-          break
-        }
-      default:
-        return
+
+  private func download() async {
+    status = .connecting
+
+    do {
+      // Connect to file transfer server (already includes +1 in serverPort from init)
+      let socket = try await NetSocketNew.connect(
+        host: serverAddress,
+        port: serverPort,
+        tls: .disabled
+      )
+      defer { Task { await socket.close() } }
+
+      status = .connected
+
+      // Send magic header
+      let headerData = Data(endian: .big) {
+        "HTXF".fourCharCode()
+        self.referenceNumber
+        UInt32.zero
+        UInt32.zero
       }
-    }
-    
-    self.status = .connecting
-    self.connection?.start(queue: .global())
-  }
-  
-  private func sendMagic() {
-    guard let c = connection, self.status == .connected else {
-      self.invalidate()
-      print("HotlineFileClient: invalid connection to send header.")
-      return
-    }
-    
-    let headerData = Data(endian: .big) {
-      "HTXF".fourCharCode()
-      self.referenceNumber
-      UInt32.zero
-      UInt32.zero
-    }
-    
-    c.send(content: headerData, completion: .contentProcessed { [weak self] (error) in
-      guard let self = self else {
-        return
+      try await socket.write(headerData)
+
+      status = .progress(0.0)
+
+      // Download file data with progress updates
+      let fileData = try await socket.read(Int(referenceDataSize)) { current, total in
+        self.status = .progress(Double(current) / Double(total))
       }
-      
-      guard error == nil else {
-        self.status = .failed(.failedToConnect)
-        self.invalidate()
-        return
+
+      print("HotlineFilePreviewClient: Complete")
+      status = .completed
+
+      // Notify delegate on main thread
+      let reference = self.referenceNumber
+      await MainActor.run {
+        self.delegate?.hotlineFilePreviewComplete(client: self, reference: reference, data: fileData)
       }
-      
-      self.status = .progress(0.0)
-      self.receive()
-    })
-  }
-  
-  private func receive() {
-    guard let c = self.connection else {
+
+    } catch is CancellationError {
+      // Already handled in cancel()
       return
-    }
-    
-    c.receive(minimumIncompleteLength: 1, maximumLength: Int(UInt16.max)) { [weak self] (data, context, isComplete, error) in
-      guard let self = self else {
-        return
-      }
-      
-      guard error == nil else {
-        self.status = .failed(.failedToDownload)
-        self.invalidate()
-        return
-      }
-      
-      if let newData = data, !newData.isEmpty {
-        self.fileBytesTransferred += newData.count
-        self.fileBytes.append(newData)
-        self.status = .progress(Double(self.fileBytesTransferred) / Double(self.referenceDataSize))
-        print("HotlineFilePreviewClient: Download progress", self.fileBytesTransferred, self.referenceDataSize, isComplete)
-      }
-      
-      if self.fileBytesTransferred < Int(self.referenceDataSize) {
-        self.receive()
-      }
-      else {
-        print("HotlineFilePreviewClient: Complete")
-        let data = self.fileBytes
-        
-        self.status = .completed
-        self.invalidate()
-        
-        let reference = self.referenceNumber
-        DispatchQueue.main.sync {
-          self.delegate?.hotlineFilePreviewComplete(client: self, reference: reference, data: data)
-        }
+    } catch {
+      print("HotlineFilePreviewClient: Download failed: \(error)")
+
+      if status == .connecting {
+        status = .failed(.failedToConnect)
+      } else {
+        status = .failed(.failedToDownload)
       }
     }
   }
diff --git a/Hotline/Library/NetSocketNew.swift b/Hotline/Library/NetSocketNew.swift
index 753a201..d4507b8 100644
--- a/Hotline/Library/NetSocketNew.swift
+++ b/Hotline/Library/NetSocketNew.swift
@@ -181,16 +181,17 @@ public actor NetSocketNew {
   
   /// Connect to a remote host and return a ready socket
   ///
-  /// This method establishes a TCP connection and waits until the connection is in `.ready` state.
+  /// This method establishes a TCP connection using Network framework types and waits until
+  /// the connection is in `.ready` state.
   ///
   /// - Parameters:
-  ///   - host: Hostname or IP address to connect to
-  ///   - port: Port number (0-65535)
+  ///   - host: Network framework host (e.g., `.name("example.com", nil)` or `.ipv4(...)`)
+  ///   - port: Network framework port
   ///   - tls: TLS policy (default: enabled with default settings)
   ///   - config: Socket configuration (default: standard settings)
   /// - Returns: A connected and ready `NetSocketNew`
-  /// - Throws: `NetSocketError.invalidPort` if port is invalid, or network errors
-  public static func connect(host: String, port: UInt16, tls: TLSPolicy = .enabled(), config: Config = .init()) async throws -> NetSocketNew {
+  /// - Throws: Network errors or connection failures
+  public static func connect(host: NWEndpoint.Host, port: NWEndpoint.Port, tls: TLSPolicy = .enabled(), config: Config = .init()) async throws -> NetSocketNew {
     let parameters = NWParameters.tcp
     if tls.enabled {
       let tlsOptions = NWProtocolTLS.Options()
@@ -198,16 +199,20 @@ public actor NetSocketNew {
       parameters.defaultProtocolStack.applicationProtocols.insert(tlsOptions, at: 0)
     }
 
-    guard let nwPort = NWEndpoint.Port(rawValue: port) else {
-      throw NetSocketError.invalidPort
-    }
-    
-    let conn = NWConnection(host: .name(host, nil), port: nwPort, using: parameters)
+    let conn = NWConnection(host: host, port: port, using: parameters)
     let socket = NetSocketNew(connection: conn, config: config)
     try await socket.start()
     return socket
   }
-  
+
+  /// Convenience wrapper to connect using string hostname and integer port
+  public static func connect(host: String, port: UInt16, tls: TLSPolicy = .enabled(), config: Config = .init()) async throws -> NetSocketNew {
+    guard let nwPort = NWEndpoint.Port(rawValue: port) else {
+      throw NetSocketError.invalidPort
+    }
+    return try await connect(host: .name(host, nil), port: nwPort, tls: tls, config: config)
+  }
+
   /// Inject custom encoding/decoding logic (supports any encoder/decoder: JSON, CBOR, MessagePack, etc.)
   ///
   /// Example with JSONEncoder:
@@ -738,7 +743,46 @@ public actor NetSocketNew {
       guard !isClosed else { throw NetSocketError.closed }
     }
   }
-  
+
+  /// Read exactly N bytes with progress callbacks
+  ///
+  /// Like `read(_:)`, but reads in chunks and reports progress after each chunk.
+  /// Useful for downloading large amounts of data where you want to update UI progress.
+  ///
+  /// Example:
+  /// ```swift
+  /// let data = try await socket.read(1_000_000) { current, total in
+  ///   print("Progress: \(current)/\(total)")
+  /// }
+  /// ```
+  ///
+  /// - Parameters:
+  ///   - count: Number of bytes to read
+  ///   - chunkSize: Size of chunks to read at a time (default: 8192)
+  ///   - progress: Optional callback with (bytesReceived, totalBytes)
+  /// - Returns: Exactly `count` bytes
+  /// - Throws: `NetSocketError` if connection closes before enough data arrives
+  public func read(
+    _ count: Int,
+    chunkSize: Int = 8192,
+    progress: (@Sendable (Int, Int) -> Void)? = nil
+  ) async throws -> Data {
+    var data = Data()
+    data.reserveCapacity(count)
+    var received = 0
+
+    while received < count {
+      try Task.checkCancellation()
+      let toRead = min(chunkSize, count - received)
+      let chunk = try await read(toRead)
+      data.append(chunk)
+      received += chunk.count
+      progress?(received, count)
+    }
+
+    return data
+  }
+
   func peek(_ count: Int) async throws -> Data {
     try await ensureReadable(count)
     let slice = buffer[head..<(head + count)]