hope-voice-api/Sources/LiveKit/VoiceSession.swift

import UIKit
import Foundation
import LiveKit
import os

private let logger = Logger(subsystem: "dev.shane.voice", category: "session")

enum VoiceSessionState: Equatable {
    case idle
    case connecting
    case waitingAgent
    case ready
    case failed(String)
}

struct AgentInfo: Equatable {
    let identity: String
    let name: String?
}

/// Room lifecycle: connect, agent discovery, agent_ready wait, reconnect-once,
/// disconnect teardown. Ports the web UI's connect flow.
@MainActor
final class VoiceSession: ObservableObject {
    @Published private(set) var state: VoiceSessionState = .idle
    @Published private(set) var agent: AgentInfo?
    @Published private(set) var isMuted = false
    @Published private(set) var micError: String?
    @Published private(set) var models: [ControlChannel.ModelInfo] = []
    @Published private(set) var selectedModel: String?

    private(set) var room = Room()
    let controlChannel = ControlChannel()
    let transcriptHandler = TranscriptHandler()
    let stateModel = ConversationStateModel()
    private let config: ServerConfig
    private let tokenClient: TokenClient
    private let settings: SettingsStore
    private var reconnectAttempted = false
    private var connectAttempts = 0
    private var agentWaitTask: Task<Void, Never>?

    init(config: ServerConfig, tokenClient: TokenClient, settings: SettingsStore) {
        self.config = config
        self.tokenClient = tokenClient
        self.settings = settings
        room.add(delegate: self)
        room.add(delegate: controlChannel)
        room.add(delegate: transcriptHandler)
        controlChannel.onEvent = { [weak self] event in
            Task { @MainActor [weak self] in
                self?.handleControlEvent(event)
            }
        }
    }

    private func handleControlEvent(_ event: ControlChannel.InboundEvent) {
        switch event {
        case .models(let received, let selected):
            models = received
            if selectedModel == nil, let selected {
                selectedModel = selected
                settings.modelID = selected
            }
            logger.log("models received count=\(received.count) selected=\(selected ?? "nil", privacy: .public)")
        case .modelChanged(let model):
            selectedModel = model
            settings.modelID = model
            logger.log("model_changed applied: \(model, privacy: .public)")
        default:
            break
        }
    }

    var isReady: Bool { state == .ready }

    func connect() async {
        await start()
    }

    private func makeRoom() -> Room {
        let room = Room()
        room.add(delegate: self)
        room.add(delegate: controlChannel)
        room.add(delegate: transcriptHandler)
        return room
    }

    private func start() async {
        state = .connecting
        connectAttempts = 0
        await attemptConnect()
    }

    private func attemptConnect() async {
        let result: Result<Void, Error>
        do {
            let token = try await tokenClient.fetchToken()
            logger.log("token ok room=\(token.room, privacy: .public)")
            try await room.connect(url: config.livekitURL.absoluteString, token: token.token, connectOptions: makeConnectOptions())
            result = .success(())
        } catch {
            result = .failure(error)
        }
        switch result {
        case .success:
            logger.log("room connected")
            UIApplication.shared.isIdleTimerDisabled = true
            scanForAgent()
            startAgentWait()
        case .failure(let error):
            let isTimeout = (error as? LiveKitError)?.type == .timedOut
            if isTimeout, connectAttempts < 2 {
                connectAttempts += 1
                let attempt = connectAttempts
                logger.log("connect attempt \(attempt) timed out, retrying with fresh room")
                await teardownRoom()
                await attemptConnect()
                return
            }
            logger.error("connect failed: \(String(describing: error), privacy: .public)")
            state = .failed(error.localizedDescription)
        }
    }

    private func makeConnectOptions() -> ConnectOptions {
        return ConnectOptions(
            socketConnectTimeoutInterval: 15,
            primaryTransportConnectTimeout: 20,
            publisherTransportConnectTimeout: 20
        )
    }

    private func teardownRoom() async {
        await room.disconnect()
        room = makeRoom()
    }

    private func handleAgentTimeout() async {
        guard state == .waitingAgent, !reconnectAttempted else {
            if state == .waitingAgent {
                state = .failed("Agent never joined. Is the agent worker running?")
            }
            return
        }
        reconnectAttempted = true
        await teardownRoom()
        await start()
    }

    private func startAgentWait() {
        state = .waitingAgent
        agentWaitTask?.cancel()
        agentWaitTask = Task { [weak self] in
            try? await Task.sleep(for: .seconds(10))
            guard !Task.isCancelled else { return }
            await self?.handleAgentTimeout()
        }
    }

    func agentReadyReceived() {
        guard state == .waitingAgent else {
            logger.log("agent_ready ignored in state \(String(describing: self.state), privacy: .public)")
            return
        }
        agentWaitTask?.cancel()
        state = .ready
        stateModel.start()
        logger.log("state=ready, enabling mic")
        Task { await enableMicrophone() }
        Task { await syncSettingsOnJoin() }
    }

    private func syncSettingsOnJoin() async {
        guard isReady else { return }
        let voiceID = settings.voiceID
        do {
            try await controlChannel.setVoice(voiceID, room: room)
            logger.log("sync-on-join set_voice sent: \(voiceID, privacy: .public)")
            if let model = selectedModel ?? settings.modelID {
                try await controlChannel.setModel(model, room: room)
                logger.log("sync-on-join set_model sent: \(model, privacy: .public)")
            }
        } catch {
            logger.error("sync-on-join failed: \(String(describing: error), privacy: .public)")
        }
    }

    func selectVoice(_ id: String) async {
        settings.voiceID = id
        guard isReady else { return }
        do {
            try await controlChannel.setVoice(id, room: room)
            logger.log("set_voice sent: \(id, privacy: .public)")
        } catch {
            logger.error("set_voice failed: \(String(describing: error), privacy: .public)")
        }
    }

    func selectModel(_ id: String) async {
        settings.modelID = id
        guard isReady else { return }
        do {
            try await controlChannel.setModel(id, room: room)
            logger.log("set_model sent: \(id, privacy: .public)")
        } catch {
            logger.error("set_model failed: \(String(describing: error), privacy: .public)")
        }
    }

    private func scanForAgent() {
        for participant in room.remoteParticipants.values {
            if isAgentParticipant(participant) {
                foundAgent(participant)
                return
            }
        }
    }

    private func isAgentParticipant(_ participant: RemoteParticipant) -> Bool {
        participant.kind == .agent || participant.identity?.stringValue == "voice-assistant"
    }

    private func foundAgent(_ participant: RemoteParticipant) {
        agent = AgentInfo(identity: participant.identity?.stringValue ?? "", name: participant.name)
    }

    func enableMicrophone() async {
        for attempt in 0..<3 {
            do {
                try await room.localParticipant.setMicrophone(enabled: true)
                isMuted = false
                micError = nil
                logger.log("mic enabled")
                return
            } catch {
                let message = String(describing: error)
                if attempt < 2 {
                    logger.log("mic enable failed (\(message, privacy: .public)), retrying in 1s")
                    try? await Task.sleep(for: .seconds(1))
                    continue
                }
                // Non-fatal: keep the session usable (text mode still works).
                logger.error("mic enable failed after retries: \(message, privacy: .public)")
                micError = "Microphone unavailable (\(error.localizedDescription)). Text mode still works."
            }
        }
    }

    func setMuted(_ muted: Bool) async {
        guard isReady else { return }
        do {
            try await room.localParticipant.setMicrophone(enabled: !muted)
            isMuted = muted
        } catch {
            micError = "Microphone error: \(error.localizedDescription)"
        }
    }

    func setMode(_ mode: String) async {
        guard isReady else { return }
        do {
            try await controlChannel.setMode(mode, room: room)
            logger.log("set_mode sent: \(mode, privacy: .public)")
        } catch {
            logger.error("set_mode failed: \(String(describing: error), privacy: .public)")
        }
    }

    func sendText(_ text: String) async throws {
        logger.log("sendText lk.chat chars=\(text.count)")
        try await room.localParticipant.sendText(text, for: "lk.chat")
    }

    func disconnect() async {
        agentWaitTask?.cancel()
        state = .idle
        agent = nil
        isMuted = false
        micError = nil
        models = []
        reconnectAttempted = false
        stateModel.stop()
        UIApplication.shared.isIdleTimerDisabled = false
        await teardownRoom()
    }

    func tickAudioLevels(local: Float, remote: Float) {
        stateModel.tick(localAudioLevel: local, remoteAudioLevel: remote)
    }

}

extension VoiceSession: RoomDelegate {
    nonisolated func room(_ room: Room, participantDidConnect participant: RemoteParticipant) {
        Task { @MainActor [weak self] in
            guard let self, self.isAgentParticipant(participant) else { return }
            self.foundAgent(participant)
        }
    }

    nonisolated func room(_ room: Room, didUpdateSpeakingParticipants participants: [Participant]) {
        Task { @MainActor [weak self] in
            guard let self, self.state == .ready else { return }
            var local = Float(0)
            var remote = Float(0)
            for participant in participants {
                if participant.identity?.stringValue == room.localParticipant.identity?.stringValue {
                    local = max(local, participant.audioLevel)
                } else {
                    remote = max(remote, participant.audioLevel)
                }
            }
            self.tickAudioLevels(local: local, remote: remote)
        }
    }

    nonisolated func room(_ room: Room, participant: RemoteParticipant?, didReceiveData data: Data, forTopic topic: String, encryptionType: EncryptionType) {
        logger.log("data topic=\(topic, privacy: .public) bytes=\(data.count) payload=\(String(data: data.prefix(300), encoding: .utf8) ?? "<binary>", privacy: .public)")
        Task { @MainActor [weak self] in
            guard let self else { return }
            // Secondary transcript source (finals only).
            if topic == "transcript",
               let json = try? JSONDecoder().decode([String: JSONValue].self, from: data),
               json["type"]?.stringValue == "transcript",
               let text = json["text"]?.stringValue {
                let role: TranscriptEvent.Role = json["role"]?.stringValue == "agent" ? .agent : .user
                self.transcriptHandler.receiveDataTranscript(role: role, text: text)
                return
            }
            if self.controlChannel.handleTaskEvent(data) { return }
            _ = self.controlChannel.handleVoiceControl(data, session: self)
        }
    }
}