pi-remote-ios/Sources/Core/Sessions/SessionConnection.swift

// SessionConnection.swift
// IC-2.1 — one WebSocket connection per session.
//
// Lifecycle:
//   init → resume(from:) → [stream data] → suspend() → resume(from:) → …
//
// URL pattern:  ws://<host>:<port>/sessions/<id>/stream?token=<bearerToken>
// (TLS pinning is wired in a follow-up task; plain `ws://` for now.)

import Combine
import Foundation

// MARK: - SessionConnection

/// Manages a single IC-1 WebSocket session stream.
///
/// Conforms to `ObservableObject` so SwiftUI views can react to
/// `connectionState` changes without manually subscribing to Combine.
///
/// All mutable state is main-actor-isolated. Callers on background contexts
/// must dispatch accordingly (normal for `@MainActor` types).
@MainActor
public final class SessionConnection: ObservableObject {

    // MARK: - Identity

    /// The session identifier used in the URL path and scrollback file name.
    public let id: String

    // MARK: - Combine publishers

    /// Emits raw ANSI bytes in arrival order (binary frames, header stripped).
    public let stream = PassthroughSubject<Data, Never>()

    /// Emits every JSON frame received from the server.
    public let stateEvents = PassthroughSubject<ServerToClient, Never>()

    /// Emits snapshot content ready to feed directly into SwiftTerm:
    /// ESC[H + ESC[2J (clear+home) followed by the pane's current content.
    public let snapshots = PassthroughSubject<Data, Never>()

    /// Tracks the WebSocket lifecycle.
    @Published public private(set) var connectionState: ConnectionState = .disconnected

    // MARK: - Scrollback

    /// Persistent rolling ANSI cache for this session.
    public private(set) var scrollback: ScrollbackCache

    // MARK: - Private

    private let credential: SidecarCredential
    private var client: WebSocketClient?
    private var cancellables = Set<AnyCancellable>()

    // MARK: - Init

    /// Creates a `SessionConnection` for `id` authenticated with `credential`.
    ///
    /// Does **not** open a WebSocket. Call `resume(from:)` to connect.
    init(id: String, credential: SidecarCredential) {
        self.id = id
        self.credential = credential
        self.scrollback = ScrollbackCache(sessionId: id)
    }

    // MARK: - Public API

    /// Opens (or re-opens) the WebSocket and sends a `resume` frame.
    ///
    /// - Parameter lastSeq: The last acknowledged sequence number, or `nil`
    ///   to request replay from the beginning.
    public func resume(from lastSeq: UInt64?) async {
        guard let url = streamURL else {
            #if DEBUG
            print("[SessionConnection] Could not construct stream URL for session \(id) — aborting resume.")
            #endif
            return
        }

        // Tear down any existing connection cleanly before reconnecting.
        await suspend()

        let ws = WebSocketClient()
        client = ws

        // Mirror WebSocketClient's connection state into our @Published property.
        ws.connectionState
            .receive(on: DispatchQueue.main)
            .sink { [weak self] state in
                self?.connectionState = state
            }
            .store(in: &cancellables)

        // Binary frames → scrollback + downstream `stream` subject.
        ws.incomingBinary
            .sink { [weak self] frame in
                guard let self else { return }
                self.scrollback.append(frame.data)
                self.stream.send(frame.data)
            }
            .store(in: &cancellables)

        // JSON frames → `stateEvents` + snapshot converter.
        ws.incomingJSON
            .sink { [weak self] frame in
                guard let self else { return }
                if case .snapshot(_, let base64) = frame {
                    // Decode base64 → text, prepend clear+home, normalise line endings.
                    if let raw = Data(base64Encoded: base64),
                       let text = String(data: raw, encoding: .utf8) {
                        let header = "\u{1B}[H\u{1B}[2J"  // cursor home + clear screen
                        let body   = text.replacingOccurrences(of: "\n", with: "\r\n")
                        self.snapshots.send(Data((header + body).utf8))
                    }
                }
                self.stateEvents.send(frame)
            }
            .store(in: &cancellables)

        // Once connected, send the appropriate opening frame.
        //
        // • lastSeq == nil  → fresh attach: live output already flows; caller
        //   will send resize then snapshotRequest. No frame sent here.
        // • lastSeq != nil  → reconnect after gap: replay missed output.
        ws.connectionState
            .filter { $0 == .connected }
            .first()
            .sink { [weak self, weak ws, lastSeq] _ in
                guard let self, let ws else { return }
                if let seq = lastSeq {
                    Task { @MainActor [self, ws, seq] in
                        try? await ws.send(.resume(lastSeq: seq))
                        _ = self
                    }
                }
                // fresh connect: caller drives resize → snapshotRequest
            }
            .store(in: &cancellables)

        ws.connect(url: url)
    }

    /// Sends a frame to the server.
    ///
    /// - Throws: `WebSocketClientError.notConnected` if there is no active
    ///   socket, or `WebSocketClientError.encodingFailed` on serialisation
    ///   failure.
    public func send(_ frame: ClientToServer) async throws {
        guard let client else {
            throw WebSocketClientError.notConnected
        }
        try await client.send(frame)
    }

    /// Closes the WebSocket but keeps local state (scrollback + cursor).
    public func suspend() async {
        client?.disconnect()
        client = nil
        cancellables.removeAll()
        connectionState = .disconnected
    }

    // MARK: - URL construction

    /// Builds `ws://<host>:<port>/sessions/<id>/stream?token=<bearerToken>`.
    ///
    /// Returns `nil` if `URLComponents` cannot produce a valid URL (should
    /// never happen in practice with well-formed credentials).
    ///
    /// Note: plain `ws://` is used for now; TLS + cert-pinning wired in
    /// the T-2.5 follow-up task.
    private var streamURL: URL? {
        var components = URLComponents()
        components.scheme = "ws"
        components.host = credential.host
        components.port = credential.port
        components.path = "/sessions/\(id)/stream"
        components.queryItems = [
            URLQueryItem(name: "token", value: credential.bearerToken)
        ]
        return components.url
    }
}