rainbow-dragon/Sources/ForgeCameraARRI/ArriDriver.swift

import Foundation
import ForgeCamera
import ForgeGrade
import ForgeColor
import ForgeOffload

#if canImport(Glibc)
import Glibc
#endif

/// ALEXA 35 / Mini LF driver — real ARRI CAP protocol:
/// binary TCP :5055, big-endian frames, MD5 challenge auth, variable subscription.
/// Verified against CAP v1.12 spec tables (docs/research/arri-cap-protocol.md).
/// CDL pushed via SetVariable 0x0050 (10×F32 blob). 3D LUT via Set3DLutData —
/// marked experimental until chunk framing verified on hardware.
public actor ArriDriver: CameraDriver {
    public nonisolated let capabilities = CameraCapabilities(
        vendor: .arri,
        supportsNativeCDL: true,
        lut3dSizes: [17, 33, 65],
        metadataFields: [.clipName, .exposureIndex, .whiteBalance, .tint, .fps, .timecode])

    private let host: String
    private let port: Int
    private let password: String
    private let clientName: String

    private var socketFD: Int32 = -1
    private var readTask: Task<Void, Never>?
    private var keepAliveTask: Task<Void, Never>?
    private var nextMsgId: UInt16 = 1

    private var readBuffer = Data()
    private var pendingReplies: [UInt16: CheckedContinuation<CAP.Frame, Error>] = [:]

    // Variable cache -> CameraState.
    private var varCache: [UInt16: Data] = [:]
    private var lastState = CameraState(connection: .disconnected)
    private var stateContinuations: [UUID: AsyncStream<CameraState>.Continuation] = [:]

    public init(host: String, port: Int = 5055, password: String = "arri", clientName: String = "Forge") {
        self.host = host
        self.port = port
        self.password = password
        self.clientName = clientName
    }

    public nonisolated var state: AsyncStream<CameraState> {
        AsyncStream { continuation in
            let id = UUID()
            Task { await self.register(id: id, continuation: continuation) }
            continuation.onTermination = { _ in
                Task { await self.unregister(id: id) }
            }
        }
    }

    private func register(id: UUID, continuation: AsyncStream<CameraState>.Continuation) {
        stateContinuations[id] = continuation
    }

    private func unregister(id: UUID) {
        stateContinuations.removeValue(forKey: id)
    }

    private func emit(_ s: CameraState) {
        lastState = s
        for c in stateContinuations.values {
            c.yield(s)
        }
    }

    // MARK: Socket

    private func openSocket() throws {
        let fd = socket(AF_INET, Int32(SOCK_STREAM.rawValue), 0)
        guard fd >= 0 else { throw CameraError.connectionFailed("socket() failed") }
        var tv = timeval(tv_sec: 3, tv_usec: 0)
        setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, socklen_t(MemoryLayout<timeval>.size))
        setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, socklen_t(MemoryLayout<timeval>.size))

        var addr = sockaddr_in()
        addr.sin_family = sa_family_t(AF_INET)
        addr.sin_port = in_port_t(UInt16(port).bigEndian)
        guard inet_pton(AF_INET, host, &addr.sin_addr) == 1 else {
            close(fd)
            throw CameraError.connectionFailed("bad host \(host)")
        }
        let rc = withUnsafePointer(to: &addr) { ptr in
            ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { sa in
                Glibc.connect(fd, sa, socklen_t(MemoryLayout<sockaddr_in>.size))
            }
        }
        guard rc == 0 else {
            close(fd)
            throw CameraError.connectionFailed("connect \(host):\(port) failed")
        }
        socketFD = fd
    }

    private func closeSocket() {
        if socketFD >= 0 {
            close(socketFD)
            socketFD = -1
        }
    }

    private func sendFrame(_ frame: CAP.Frame) throws {
        guard socketFD >= 0 else { throw CameraError.connectionFailed("not connected") }
        let data = CAP.encode(frame)
        var sent = 0
        try data.withUnsafeBytes { (buf: UnsafeRawBufferPointer) in
            while sent < buf.count {
                let n = write(socketFD, buf.baseAddress!.advanced(by: sent), buf.count - sent)
                guard n > 0 else { throw CameraError.connectionFailed("write failed") }
                sent += n
            }
        }
    }

    /// Send command, await matching reply (msgId correlation).
    private func request(_ cmd: CAP.Command, payload: Data = Data(), timeout: Duration = .seconds(3)) async throws -> CAP.Frame {
        let msgId = nextMsgId
        nextMsgId &+= 1
        if nextMsgId == 0 { nextMsgId = 1 }

        let frame = CAP.Frame(
            msgType: CAP.MsgType.command.rawValue,
            msgId: msgId,
            cmdCode: cmd.rawValue,
            payload: payload)

        return try await withCheckedThrowingContinuation { cont in
            pendingReplies[msgId] = cont
            do {
                try sendFrame(frame)
            } catch {
                pendingReplies.removeValue(forKey: msgId)
                cont.resume(throwing: error)
                return
            }
            Task {
                try? await Task.sleep(for: timeout)
                self.timeoutReply(msgId: msgId)
            }
        }
    }

    private func timeoutReply(msgId: UInt16) {
        if let cont = pendingReplies.removeValue(forKey: msgId) {
            cont.resume(throwing: CameraError.connectionFailed("reply timeout msgId \(msgId)"))
        }
    }

    // MARK: Read loop

    private func startReadLoop() {
        readTask?.cancel()
        readTask = Task.detached { [weak self] in
            guard let self else { return }
            let fd = await self.socketFD
            var chunk = [UInt8](repeating: 0, count: 64 * 1024)
            while !Task.isCancelled {
                let n = read(fd, &chunk, chunk.count)
                if n > 0 {
                    await self.ingest(Data(chunk[0..<n]))
                } else if n == 0 {
                    await self.handleDisconnect()
                    return
                } else {
                    // EAGAIN on timeout: keep looping unless cancelled.
                    if errno != EAGAIN && errno != EWOULDBLOCK {
                        await self.handleDisconnect()
                        return
                    }
                }
            }
        }
    }

    private func ingest(_ data: Data) {
        readBuffer.append(data)
        while let frame = try? CAP.decodeFirst(&readBuffer) {
            route(frame)
        }
    }

    private func route(_ frame: CAP.Frame) {
        switch frame.msgType {
        case CAP.MsgType.reply.rawValue:
            if let cont = pendingReplies.removeValue(forKey: frame.msgId) {
                cont.resume(returning: frame)
            } else if frame.cmdCode == CAP.Command.welcome.rawValue || frame.msgId == 0 {
                // Unsolicited Welcome on connect — stored implicitly; nothing pending.
            }
        case CAP.MsgType.event.rawValue:
            // Variable update: payload = U16 id + value.
            if let id = CAP.readU16(frame.payload, at: 0) {
                varCache[id] = frame.payload.subdata(in: frame.payload.startIndex + 2..<frame.payload.endIndex)
                rebuildState()
            }
        default:
            break
        }
    }

    private func handleDisconnect() {
        closeSocket()
        for (_, cont) in pendingReplies {
            cont.resume(throwing: CameraError.connectionFailed("connection closed"))
        }
        pendingReplies.removeAll()
        if lastState.connection == .connected {
            emit(CameraState(connection: .disconnected))
        }
    }

    // MARK: State assembly

    private func rebuildState() {
        func u16(_ v: CAP.Variable) -> UInt16? {
            varCache[v.rawValue].flatMap { CAP.readU16($0, at: 0) }
        }
        func u32(_ v: CAP.Variable) -> UInt32? {
            varCache[v.rawValue].flatMap { CAP.readU32($0, at: 0) }
        }
        func f32(_ v: CAP.Variable) -> Float? {
            varCache[v.rawValue].flatMap { CAP.readF32($0, at: 0) }
        }

        let bits = CAP.CameraStateBits(rawValue: u16(.cameraState) ?? 0)
        // Clip name composed from reel + clip number (full names need GetClipList).
        var clipName: String?
        if let reel = u16(.currentReel), let clip = u16(.clipNumber), clip > 0 {
            clipName = String(format: "R%03dC%03d", reel, clip)
        }
        let newState = CameraState(
            connection: .connected,
            isRecording: bits.contains(.recording),
            clipName: clipName,
            metadata: CameraMetadata(
                exposureIndex: u32(.exposureIndex).map(Int.init),
                whiteBalance: u32(.colorTemperature).map(Int.init),
                tint: f32(.tint).map { Int($0.rounded()) },
                fps: f32(.sensorFPS).map(Double.init),
                timecode: u32(.timecode).map(CAP.decodeTimecodeBCD)))
        if newState != lastState {
            emit(newState)
        }
    }

    // MARK: CameraDriver

    public func connect() async throws {
        closeSocket()
        readBuffer.removeAll()
        varCache.removeAll()
        try openSocket()
        startReadLoop()

        // Camera sends Welcome unsolicited; brief settle for it to arrive.
        try? await Task.sleep(for: .milliseconds(50))

        // Auth: challenge -> MD5(challenge + password) -> client name.
        let challengeReply = try await request(.requestPwdChallenge)
        guard challengeReply.cmdCode == CAP.Result.ok.rawValue,
              let (challenge, _) = CAP.readString(challengeReply.payload, at: 0) else {
            await teardownAfterFailure()
            throw CameraError.connectionFailed("challenge request failed")
        }

        let hash = MD5.hexString(Data((challenge + password).utf8))
        var pwPayload = Data()
        CAP.putString(hash, into: &pwPayload)
        let pwReply = try await request(.password, payload: pwPayload)
        guard pwReply.cmdCode == CAP.Result.ok.rawValue else {
            await teardownAfterFailure()
            throw CameraError.connectionFailed("password rejected (result \(pwReply.cmdCode))")
        }

        var namePayload = Data()
        CAP.putString(clientName, into: &namePayload)
        _ = try await request(.clientName, payload: namePayload)

        // Subscribe to state + metadata variables — EVENT-driven from here.
        var subPayload = Data()
        for v in [CAP.Variable.cameraState, .currentReel, .clipNumber, .exposureIndex,
                  .colorTemperature, .tint, .sensorFPS, .timecode] {
            CAP.putU16(v.rawValue, into: &subPayload)
        }
        _ = try await request(.requestVariables, payload: subPayload)

        emit(CameraState(connection: .connected))
        startKeepAlive()
    }

    private func teardownAfterFailure() async {
        readTask?.cancel()
        readTask = nil
        closeSocket()
    }

    public func disconnect() async {
        keepAliveTask?.cancel()
        keepAliveTask = nil
        readTask?.cancel()
        readTask = nil
        closeSocket()
        emit(CameraState(connection: .disconnected))
    }

    /// CAP requires client traffic every 1 s.
    private func startKeepAlive() {
        keepAliveTask?.cancel()
        keepAliveTask = Task {
            while !Task.isCancelled {
                try? await Task.sleep(for: .milliseconds(900))
                _ = try? await self.request(.live, timeout: .seconds(2))
            }
        }
    }

    public func push(look: FlattenedLook) async throws {
        if let cdl = look.cdl {
            var payload = Data()
            CAP.putU16(CAP.Variable.cdlValues.rawValue, into: &payload)
            payload.append(CAP.encodeCDLBlob(
                slope: (cdl.slope.x, cdl.slope.y, cdl.slope.z),
                offset: (cdl.offset.x, cdl.offset.y, cdl.offset.z),
                power: (cdl.power.x, cdl.power.y, cdl.power.z),
                saturation: cdl.saturation))
            let reply = try await request(.setVariable, payload: payload)
            guard reply.cmdCode == CAP.Result.ok.rawValue else {
                throw CameraError.pushFailed("SetVariable CDL -> result \(reply.cmdCode)")
            }
        }
        if let lut = look.lut, !lut.isApproximatelyIdentity() {
            guard capabilities.lut3dSizes.contains(lut.size) else {
                throw CameraError.unsupportedLook("LUT size \(lut.size) not in \(capabilities.lut3dSizes)")
            }
            // EXPERIMENTAL: Set3DLutData payload layout unverified on hardware
            // (chunk framing sections unavailable publicly). Sim mirrors this
            // simple layout: U16 lattice size + F32 table.
            var payload = Data()
            CAP.putU16(UInt16(lut.size), into: &payload)
            for v in lut.table {
                CAP.putF32(v, into: &payload)
            }
            let reply = try await request(.set3DLutData, payload: payload, timeout: .seconds(10))
            guard reply.cmdCode == CAP.Result.ok.rawValue else {
                throw CameraError.pushFailed("Set3DLutData -> result \(reply.cmdCode)")
            }
        }
        if look.cdl == nil && (look.lut == nil || look.lut!.isApproximatelyIdentity()) {
            throw CameraError.unsupportedLook("nothing to push")
        }
    }
}