pi-remote-ios/Sources/Core/Network/FrameCodec.swift

// FrameCodec.swift
// IC-1 WebSocket frame encoding/decoding.
//
// Wire format
// -----------
// Binary frame (server → client): [seq: 8 bytes big-endian UInt64][raw ANSI bytes]
// Text frame  (both directions) : JSON, type-discriminated by the "type" key.
//
// This file has no Starscream import — pure Foundation only, fully unit-testable.

import Foundation

// MARK: - PiState

/// The reported state of the pi agent.
public enum PiState: String, Decodable, Sendable {
    case thinking
    case tool
    case idle
    case awaitingInput = "awaiting-input"
}

// MARK: - ServerToClient

/// JSON text frames sent from the server to the client.
public enum ServerToClient: Sendable {
    /// Agent state update (thinking / tool / idle / awaiting-input).
    case state(value: PiState, tool: String?, ts: Int)
    /// Full-screen snapshot encoded as base64 ANSI bytes.
    case snapshot(seq: UInt64, data: String)
    /// Session metadata pushed on connect / rename.
    case sessionMeta(name: String, description: String?, createdAt: String)
    /// Protocol-level error from the server.
    case error(code: String, message: String)
}

extension ServerToClient: Decodable {
    private enum TypeKey: String, Decodable {
        case state
        case snapshot
        case sessionMeta = "session-meta"
        case error
    }

    private enum CodingKeys: String, CodingKey {
        case type
        case value, tool, ts          // state
        case seq, data                // snapshot
        case name, description, createdAt  // session-meta
        case code, message            // error
    }

    public init(from decoder: any Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        let typeKey = try container.decode(TypeKey.self, forKey: .type)
        switch typeKey {
        case .state:
            let value = try container.decode(PiState.self, forKey: .value)
            let tool  = try container.decodeIfPresent(String.self, forKey: .tool)
            let ts    = try container.decode(Int.self, forKey: .ts)
            self = .state(value: value, tool: tool, ts: ts)

        case .snapshot:
            let seq  = try container.decode(UInt64.self, forKey: .seq)
            let data = try container.decode(String.self, forKey: .data)
            self = .snapshot(seq: seq, data: data)

        case .sessionMeta:
            let name        = try container.decode(String.self, forKey: .name)
            let description = try container.decodeIfPresent(String.self, forKey: .description)
            let createdAt   = try container.decode(String.self, forKey: .createdAt)
            self = .sessionMeta(name: name, description: description, createdAt: createdAt)

        case .error:
            let code    = try container.decode(String.self, forKey: .code)
            let message = try container.decode(String.self, forKey: .message)
            self = .error(code: code, message: message)
        }
    }
}

// MARK: - ClientToServer

/// JSON text frames sent from the client to the server.
public enum ClientToServer: Sendable {
    /// Attach / resume from the given sequence number (nil = start fresh).
    case resume(lastSeq: UInt64?)
    /// Named key press (e.g. "escape", "up", "enter").
    case key(name: String)
    /// Literal text, delivered via send-keys -l.
    case keys(data: String)
    /// Text wrapped in bracketed-paste sequences.
    case paste(data: String)
    /// Request a full ANSI snapshot of the current pane.
    case snapshotRequest
}

extension ClientToServer: Encodable {
    private enum CodingKeys: String, CodingKey {
        case type, lastSeq, name, data
    }

    public func encode(to encoder: any Encoder) throws {
        var c = encoder.container(keyedBy: CodingKeys.self)
        switch self {
        case .resume(let lastSeq):
            try c.encode("resume", forKey: .type)
            // Encode null explicitly when lastSeq is nil so the server
            // receives {"type":"resume","lastSeq":null} as the spec requires.
            try c.encode(lastSeq, forKey: .lastSeq)

        case .key(let name):
            try c.encode("key", forKey: .type)
            try c.encode(name, forKey: .name)

        case .keys(let data):
            try c.encode("keys", forKey: .type)
            try c.encode(data, forKey: .data)

        case .paste(let data):
            try c.encode("paste", forKey: .type)
            try c.encode(data, forKey: .data)

        case .snapshotRequest:
            try c.encode("snapshot-request", forKey: .type)
        }
    }
}

// MARK: - BinaryFrame

/// A decoded IC-1 binary frame.
///
/// Wire layout: `[seq: 8 bytes big-endian UInt64][raw ANSI bytes …]`
public struct BinaryFrame: Sendable {
    /// Monotonically increasing chunk sequence number assigned by the server.
    public let seq: UInt64
    /// Raw ANSI terminal bytes for this chunk.
    public let data: Data

    /// Decodes a raw WebSocket binary message into a `BinaryFrame`.
    ///
    /// Returns `nil` when `raw` is shorter than the 8-byte header.
    public static func decode(_ raw: Data) -> BinaryFrame? {
        guard raw.count >= 8 else { return nil }

        // Read 8-byte big-endian UInt64 from the leading bytes.
        let seqBytes = raw.prefix(8)
        let seq = seqBytes.reduce(into: UInt64(0)) { acc, byte in
            acc = (acc << 8) | UInt64(byte)
        }

        let payload = raw.dropFirst(8)
        return BinaryFrame(seq: seq, data: Data(payload))
    }
}

// MARK: - FrameCodec (static utilities)

/// Namespace for IC-1 frame encode/decode helpers.
public enum FrameCodec {

    // Shared encoder / decoder — both are safe to use from a single actor.
    private static let encoder: JSONEncoder = {
        let enc = JSONEncoder()
        enc.outputFormatting = []   // compact, stable output
        return enc
    }()

    private static let decoder = JSONDecoder()

    /// Encodes a `ClientToServer` frame to a compact JSON string.
    ///
    /// Throws if encoding fails (in practice this should never happen for
    /// the well-typed cases defined above).
    public static func encode(_ frame: ClientToServer) throws -> String {
        let data = try encoder.encode(frame)
        // JSON serialisation always produces valid UTF-8; the force-unwrap
        // is safe and intentional — a nil result here would be a Foundation bug.
        return String(data: data, encoding: .utf8)! // swiftlint:disable:this force_unwrap
    }

    /// Decodes a JSON string into a `ServerToClient` frame.
    ///
    /// Throws `DecodingError` for malformed or unknown payloads.
    public static func decode(_ text: String) throws -> ServerToClient {
        let data = Data(text.utf8)
        return try decoder.decode(ServerToClient.self, from: data)
    }
}