voice-cat/clients/apple/iOS/VoiceCatiOS/SessionState.swift

import Foundation
import AVFoundation
import VoiceCatCore

// MARK: - Helper types

struct ChatMessage: Identifiable {
    let id = UUID()
    let timestamp: Date
    let senderName: String
    let text: String
    let scope: VoiceCatTextScope
}

struct ActivityEntry: Identifiable {
    let id = UUID()
    let timestamp: Date
    let text: String
}

struct VoiceState {
    var micActive = false
    var selfMuted = false
    var selfDeafened = false
    var serverMuted = false
    var serverDeafened = false
    var inputMode: VoiceCatInputMode = .voiceActivation
    var vadThreshold: Float = 0.025
    var level: Float = 0.0
    var currentDeviceId: String?
    var localStreamId: UInt32 = 0
    var screenSharing = false
    var screenStreamId: UInt32 = 0
}

// MARK: - SessionState

@Observable
@MainActor
final class SessionState {
    let client: VoiceCatClient
    let selfUserId: UInt32

    var channels: [Channel] = []
    var users: [User] = []
    var currentChannelId: UInt32 = 0
    var messages: [ChatMessage] = []
    var activityLog: [ActivityEntry] = []
    var voiceState = VoiceState()
    var permissions: Permissions
    var accounts: [Account] = []
    var devices: [Device] = []

    /// Host side of iOS screen-audio sharing — drains the broadcast extension's App Group ring
    /// and feeds the SCREEN_AUDIO stream this session owns. See BroadcastAudioPump.
    private let broadcastPump = BroadcastAudioPump()

    init(client: VoiceCatClient, selfUserId: UInt32, permissions: Permissions) {
        self.client = client
        self.selfUserId = selfUserId
        self.permissions = permissions
        AudioSessionManager.shared.client = client
        refreshChannels()
        refreshUsers()
        syncSelfChannel()
        refreshDevices()
        client.onEvent = { [weak self] ev in
            Task { @MainActor [weak self] in self?.handleEvent(ev) }
        }
        client.onLevel = { [weak self] _, rms in
            Task { @MainActor [weak self] in self?.voiceState.level = rms }
        }
        broadcastPump.onBroadcastStarted = { [weak self] in self?.startScreenShare() }
        broadcastPump.onBroadcastFinished = { [weak self] in self?.stopScreenShare() }
        broadcastPump.start()
        // When IOSAudioRouter changes the audio config, restart the voice path if needed so the
        // native VPIO engine (AEC/NS/AGC) engages or disengages to match the new preset/route.
        AudioSessionManager.shared.reconcileVoicePath = { [weak self] in self?.reconcileVoicePath() }
    }

    deinit {
        broadcastPump.stop()
        MainActor.assumeIsolated {
            AudioSessionManager.shared.client = nil
        }
    }

    // MARK: - Event dispatch

    func handleEvent(_ ev: VoiceCatEvent) {
        switch ev.type {
        case .channelList:
            refreshChannels()
            syncSelfChannel()
        case .userJoined, .userLeft:
            refreshUsers()
            syncSelfChannel()
        case .userUpdated:
            refreshUsers()
            syncSelfChannel()
            if let me = users.first(where: { $0.id == selfUserId }) {
                applyServerMuteState(muted: me.serverMuted, deafened: me.serverDeafened)
            }
        case .textMessage:
            let sender = users.first(where: { $0.id == ev.userId })?.nickname ?? "Unknown"
            messages.append(ChatMessage(
                timestamp: Date(timeIntervalSince1970: Double(ev.timestampUnixMs) / 1000),
                senderName: sender,
                text: ev.text ?? "",
                scope: ev.textScope))
        case .talkState:
            let talking = ev.u32a != 0
            let who = users.first(where: { $0.id == ev.userId })?.nickname ?? "user \(ev.userId)"
            addActivity(talking ? "\(who) started talking" : "\(who) stopped talking")
        case .streamStarted:
            // Our own SCREEN_AUDIO stream is live — begin draining the broadcast ring into it,
            // in the stream's effective channel mode (downmix to mono if the channel is mono).
            if ev.userId == selfUserId && ev.streamId == voiceState.screenStreamId {
                let sid = voiceState.screenStreamId
                let (r, cfg) = client.getStreamAudioConfig(userId: selfUserId, streamId: sid)
                let channels: UInt32 = (r == .ok && cfg?.stereo == true) ? 2 : 1
                let c = client
                broadcastPump.beginFeeding(streamChannels: channels) { pcm, samples, ch in
                    c.feedPcm(streamId: sid, pcm: pcm, samplesPerChannel: samples, channels: ch)
                }
                addActivity("Sharing screen audio (\(channels == 2 ? "stereo" : "mono"))")
                break
            }
            // A remote user started a stream — ensure the audio session is active so we can
            // hear them even if we haven't joined voice ourselves.
            if ev.userId != selfUserId {
                do {
                    try AudioSessionManager.shared.ensureSessionActive()
                } catch {
                    addActivity("Audio session activate failed: \(error)")
                }
            }
            AudioSessionManager.shared.logSessionState("stream started (user \(ev.userId))")
            addActivity("Stream started (user \(ev.userId))")
        case .streamStopped:
            addActivity("Stream stopped (user \(ev.userId))")
        case .joinResult:
            if ev.result == .ok {
                currentChannelId = ev.channelId
                addActivity("Joined channel")
                refreshUsers()
            } else {
                addActivity("Join failed: \(ev.result.description)")
            }
        case .error:
            addActivity("Error: \(ev.text ?? ev.result.description)")
        case .genericResult:
            if ev.result != .ok {
                addActivity("Operation failed: \(ev.result.description)")
            }
        case .accountList:
            accounts = client.listAccounts()
        default:
            break
        }
    }

    private func addActivity(_ text: String) {
        activityLog.append(ActivityEntry(timestamp: Date(), text: text))
        if activityLog.count > 500 { activityLog.removeFirst() }
    }

    // MARK: - Self-channel / server-mute sync

    /// Sync currentChannelId from the self user's channelId in the user list. Mirrors macOS
    /// MainWindowController.swift:461,491,522. The server auto-places every authed user into
    /// the Lobby (channel 1) on connect, but without this sync currentChannelId stays 0 and
    /// the mic button (gated on currentChannelId == 0) stays permanently dimmed.
    private func syncSelfChannel() {
        if let me = users.first(where: { $0.id == selfUserId }) {
            currentChannelId = me.channelId
        }
    }

    /// Apply server-side mute/deafen state — mirrors macOS MainWindowController.swift:693-700.
    /// iOS was previously ignoring server mute/deafen entirely.
    private func applyServerMuteState(muted: Bool, deafened: Bool) {
        if muted && !voiceState.serverMuted { addActivity("You have been server-muted") }
        if deafened && !voiceState.serverDeafened { addActivity("You have been server-deafened") }
        if !muted && voiceState.serverMuted { addActivity("Server mute cleared") }
        if !deafened && voiceState.serverDeafened { addActivity("Server deafen cleared") }
        voiceState.serverMuted = muted
        voiceState.serverDeafened = deafened
    }

    // MARK: - Data refresh

    func refreshChannels() { channels = client.listChannels() }
    func refreshUsers()    { users    = client.listUsers() }
    func refreshDevices()  { devices  = client.listDevices(.input) }

    // MARK: - Voice controls

    func joinChannel(_ channelId: UInt32, password: String = "") {
        client.joinChannel(channelId, password: password.isEmpty ? nil : password)
    }

    func leaveChannel() {
        client.leaveChannel()
        currentChannelId = 0
    }

    func startMicStream() {
        AVAudioApplication.requestRecordPermission { [weak self] granted in
            DispatchQueue.main.async {
                guard let self else { return }
                if granted {
                    self.doStartMicStream()
                } else {
                    self.addActivity("Microphone permission denied — grant in Settings > Privacy > Microphone")
                }
            }
        }
    }

    private func doStartMicStream() {
        do {
            try AudioSessionManager.shared.ensureSessionActive()
        } catch {
            addActivity("AVAudioSession activate failed: \(error)")
            return
        }

        // VPIO path: on the AEC presets, the native AVAudioEngine does AEC/NS/AGC and the core
        // runs in external mode (no hardware mic/playback). The mic stream is started with
        // externalFeed so the core skips the hardware capture device; setExternalPlayback makes
        // it skip the hardware playback device and deliver the mix to IOSVoiceProcessingEngine.
        //
        // ORDER MATTERS: set the external-playback flag now, but defer audioRestart() until
        // AFTER startStream (below) so the MIC LocalStream — which carries external_feed=true —
        // already exists when ensure_audio_running() derives external_capture. Restarting before
        // the stream exists makes the core reopen a hardware capture device that is never dropped
        // (the announce-result restart early-returns because the engine is already running); that
        // lingering miniaudio capture unit then fights the AVAudioEngine VPIO unit on the same
        // .voiceChat session and silences VPIO playback.
        let useVPIO = IOSAudioRouter.shared.currentConfigUsesVoiceProcessing
        client.setExternalPlayback(useVPIO)

        let desc = StreamDescriptor(kind: .mic, deviceId: voiceState.currentDeviceId, label: "Mic",
                                    externalFeed: useVPIO)
        let (result, streamId) = client.startStream(desc)
        if result == .ok {
            voiceState.micActive = true
            voiceState.localStreamId = streamId
            // Publish the active mic stream ID so IOSAudioRouter can reset the core's capture
            // channel count when the user switches mono↔stereo (selectCaptureChannels /
            // applyPreset). Without this, switching stereo→mono leaves the LocalStream's
            // capture_channels field at 2 and the next engine start still opens stereo.
            AudioSessionManager.shared.activeMicStreamId = streamId
            // Store the user's capture channel selection before the server acknowledges
            // the stream. The engine hasn't started yet at this point (it starts when
            // handle_stream_announce_result fires), so vc_set_capture_channels just
            // stores the value — no restart. ensure_audio_running() picks it up when
            // the stream is confirmed and opens the device with the right channel count.
            let channels = IOSAudioRouter.shared.captureChannels.channelCount
            if channels != 1 {
                client.setCaptureChannels(streamId: streamId, channels: channels)
            }
            if useVPIO {
                // The external-feed MIC stream now exists, so restart the core into full
                // external mode (no hardware capture/playback, mixer-timer only) — mic and
                // speaker are owned entirely by the VPIO engine, which we start right after.
                client.audioRestart()
                IOSVoiceProcessingEngine.shared.start(
                    client: client, micStreamId: streamId, captureChannels: channels)
            }
        } else {
            addActivity("Failed to start mic: \(result.description)")
            if useVPIO {  // revert external-playback mode so remote audio still plays
                client.setExternalPlayback(false)
                client.audioRestart()
            }
        }
    }

    func stopMicStream() {
        // Tear down the VPIO engine first (removes the mic tap + unregisters the mixed sink),
        // then stop the mic stream, then restore the core's hardware playback for any remaining
        // remote audio. Order matters: the mic stream must be gone before audioRestart so the
        // core opens a normal playback device (and no capture device — there's no mic stream).
        let wasVPIO = IOSVoiceProcessingEngine.shared.isRunning
        if wasVPIO {
            IOSVoiceProcessingEngine.shared.stop()
        }
        if voiceState.localStreamId != 0 {
            client.stopStream(voiceState.localStreamId)
            voiceState.localStreamId = 0
            AudioSessionManager.shared.activeMicStreamId = nil
        }
        if wasVPIO {
            client.setExternalPlayback(false)
            client.audioRestart()  // reopen hardware playback (no mic stream → no hw capture)
        }
        voiceState.micActive = false
        voiceState.level = 0
        // Do NOT deactivate the AVAudioSession here — the user may still want to hear
        // remote audio (other people talking). The session is deactivated only when
        // disconnecting from the server (see AppState.disconnect / .disconnected event).
    }

    /// Restart the voice path when the audio config changes mid-call (driven by IOSAudioRouter).
    /// If VPIO is involved on either the current or desired side, restart the mic so the native
    /// voice-processing engine engages/disengages and re-binds to the new route. Pure miniaudio
    /// config tweaks need no restart — the core's own audioRestart (already issued) handles them.
    private func reconcileVoicePath() {
        guard voiceState.micActive else { return }
        let want = IOSAudioRouter.shared.currentConfigUsesVoiceProcessing
        let have = IOSVoiceProcessingEngine.shared.isRunning
        guard want || have else { return }
        stopMicStream()
        doStartMicStream()
    }

    // MARK: - Screen audio share

    /// Called when the broadcast extension becomes active. Announces the SCREEN_AUDIO stream;
    /// feeding begins on the resulting `.streamStarted` event (see handleEvent). The actual
    /// system-audio capture happens in the ReplayKit upload extension (a separate process).
    private func startScreenShare() {
        guard voiceState.screenStreamId == 0 else { return }
        guard currentChannelId != 0 else {
            addActivity("Screen audio ignored — join a channel first")
            return
        }
        let (result, streamId) = client.startStream(
            StreamDescriptor(kind: .screenAudio, deviceId: nil, label: "Screen audio"))
        if result == .ok {
            voiceState.screenStreamId = streamId
            voiceState.screenSharing = true
            addActivity("Screen audio share starting…")
        } else {
            addActivity("Failed to start screen audio: \(result.description)")
        }
    }

    /// Called when the broadcast ends (or on disconnect). Stops feeding and the stream.
    private func stopScreenShare() {
        broadcastPump.endFeeding()
        if voiceState.screenStreamId != 0 {
            client.stopStream(voiceState.screenStreamId)
            voiceState.screenStreamId = 0
        }
        if voiceState.screenSharing {
            voiceState.screenSharing = false
            addActivity("Stopped sharing screen audio")
        }
    }

    func setMute(_ muted: Bool, deafened: Bool) {
        client.setSelfMute(micMuted: muted, deafened: deafened)
        voiceState.selfMuted = muted
        voiceState.selfDeafened = deafened
    }

    func setInputMode(_ mode: VoiceCatInputMode) {
        client.setInputMode(mode)
        voiceState.inputMode = mode
    }

    func setVadThreshold(_ threshold: Float) {
        client.setVadThreshold(threshold)
        voiceState.vadThreshold = threshold
    }

    func setPushToTalk(_ active: Bool) {
        client.setPushToTalk(active)
    }

    // MARK: - Text

    func sendText(_ text: String, scope: VoiceCatTextScope, targetId: UInt32 = 0) {
        client.sendText(scope: scope, targetId: targetId, text: text)
    }

    // MARK: - Admin

    func kickUser(_ userId: UInt32, reason: String) {
        client.kickUser(userId, reason: reason.isEmpty ? nil : reason)
    }

    func banUser(_ userId: UInt32, reason: String, expiresUnixMs: UInt64) {
        client.banUser(userId, reason: reason.isEmpty ? nil : reason, expiresUnixMs: expiresUnixMs)
    }

    func moveUser(_ userId: UInt32, toChannel channelId: UInt32) {
        client.moveUser(userId, toChannel: channelId)
    }

    func setPermissions(_ userId: UInt32, perms: Permissions) {
        client.setPermission(userId, perms: perms)
    }

    func setServerMute(_ userId: UInt32, muted: Bool, deafened: Bool) {
        client.setServerMute(userId, muted: muted, deafened: deafened)
    }

    func createChannel(_ info: ChannelEdit) {
        client.createChannel(info)
    }

    func editChannel(_ info: ChannelEdit) {
        client.editChannel(info)
    }

    func deleteChannel(_ channelId: UInt32) {
        client.deleteChannel(channelId)
    }

    func fetchAccountList() {
        client.requestAccountList()
    }

    func createAccount(username: String, password: String) {
        client.createAccount(username, password: password)
    }

    func deleteAccount(username: String) {
        client.deleteAccount(username)
    }

    func resetPassword(username: String, newPassword: String) {
        client.resetPassword(username, newPassword: newPassword)
    }
}