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 voiceSubscribed = false
    var selfMuted = false
    var selfDeafened = false
    var serverMuted = false
    var serverDeafened = false
    var inputMode: VoiceCatInputMode = .voiceActivation
    var vadThreshold: Float = 0.025
    var inputGain: Float = 1.0
    var inputNoiseReduction: Bool = false
    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] = []

    /// Back-reference to the app state. Once `SessionState.init` overwrites `client.onEvent`,
    /// `AppState.handleConnectEvent` no longer receives per-session events — so the
    /// `.disconnected` event for a LIVE session arrives here in `handleEvent`, not in AppState.
    /// This weak ref lets us hand the disconnect back to AppState (which owns the reconnect
    /// state machine) so the auto-reconnect fires. Set by AppState on auth success.
    weak var appState: AppState?

    // MARK: - Reconnect restore state
    //
    // When the iOS client auto-reconnects after a network drop, AppState captures the prior
    // session's channel + voice/mic state and asks the new SessionState (created on auth success)
    // to restore it. We rejoin the channel explicitly (the server auto-placed us in Lobby on
    // auth), and on the resulting `.joinResult` we re-arm voice subscription + mute/deafen. The
    // drive is here, not in AppState, because once SessionState is created it owns
    // `client.onEvent` and AppState no longer sees per-session events.
    private struct RestoreRequest {
        let channelId: UInt32
        let voiceSubscribed: Bool
        let micMuted: Bool
        let deafened: Bool
    }
    private var pendingRestore: RestoreRequest?
    private var didIssueRestoreJoin = false

    /// 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
        loadAndApplyVoiceSettings()
        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()
    }

    deinit {
        broadcastPump.stop()
    }

    // MARK: - Event dispatch

    func handleEvent(_ ev: VoiceCatEvent) {
        switch ev.type {
        case .channelList:
            refreshChannels()
            syncSelfChannel()
        case .userJoined:
            // ev.text = nickname, ev.channelId = the channel they joined (per voicecat.h).
            if ev.userId != selfUserId && ev.channelId == currentChannelId {
                EventFeedback.shared.play(.channelJoin)
                EventFeedback.shared.speak("\(ev.text ?? "Someone") joined")
            }
            refreshUsers()
            syncSelfChannel()
        case .userLeft:
            // Capture the leaving user's prior nickname/channel before refreshUsers() drops them.
            if ev.userId != selfUserId,
               let gone = users.first(where: { $0.id == ev.userId }),
               gone.channelId == currentChannelId {
                EventFeedback.shared.play(.channelLeave)
                EventFeedback.shared.speak("\(gone.nickname) left")
            }
            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"
            let body = ev.text ?? ""
            let isSelf = ev.userId == selfUserId
            let isPrivate = ev.textScope == .private
            messages.append(ChatMessage(
                timestamp: Date(timeIntervalSince1970: Double(ev.timestampUnixMs) / 1000),
                senderName: sender,
                text: body,
                scope: ev.textScope))
            EventFeedback.shared.play(isPrivate
                ? (isSelf ? .pmSent : .pmRecv)
                : (isSelf ? .channelSent : .channelRecv))
            if !isSelf {
                EventFeedback.shared.speak(isPrivate
                    ? "Private message from \(sender): \(body)"
                    : "\(sender): \(body)")
            }
        case .talkState:
            let talking = ev.u32a != 0
            if ev.userId == selfUserId {
                EventFeedback.shared.play(talking ? .vaStart : .vaStop)
            }
            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
            }
            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:
            if ev.userId == selfUserId {
                if ev.streamId == voiceState.localStreamId {
                    voiceState.localStreamId = 0
                    voiceState.micActive = false
                    voiceState.level = 0
                }
            } else {
                addActivity("Stream stopped (user \(ev.userId))")
            }
        case .voiceState:
            let subscribed = ev.u32a != 0
            voiceState.voiceSubscribed = subscribed
            if subscribed {
                doStartMicStream()
            } else {
                voiceState.micActive = false
                voiceState.level = 0
                EventFeedback.shared.play(.voiceOff)
            }
        case .joinResult:
            if ev.result == .ok {
                currentChannelId = ev.channelId
                addActivity("Joined channel")
                refreshUsers()
                // Reconnect restore: this was our restore-join. Now that the server has
                // processed the channel move, re-arm voice subscription (if the user was
                // transmitting before the drop) and re-apply the local mute/deafen state.
                // The server returns ok even when joining the channel we're already in, so
                // this fires reliably for the Lobby-too case.
                if didIssueRestoreJoin, let r = pendingRestore, r.channelId == ev.channelId {
                    didIssueRestoreJoin = false
                    completeRestore()
                }
            } else {
                addActivity("Join failed: \(ev.result.description)")
                // Restore-join failed (channel was deleted, became password-protected or
                // full while we were away). Give up on the voice/mute restore cleanly so we
                // don't leave dangling state or attempt voice without being in a channel.
                if didIssueRestoreJoin {
                    didIssueRestoreJoin = false
                    pendingRestore = nil
                }
            }
        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()
        case .disconnected:
            // Audible cue, then hand the disconnect back to AppState so its reconnect state
            // machine fires. This is the ONLY way AppState learns a live session dropped —
            // after auth success, `SessionState.init` overwrites `client.onEvent`, so
            // `AppState.handleConnectEvent` never sees this event. (Without this callback, a
            // network drop on a live session would just play the cue and leave the session as a
            // zombie — the user would have to tap Disconnect manually.)
            EventFeedback.shared.play(ev.result == .ok ? .logout : .connectionLost)
            EventFeedback.shared.speak(ev.result == .ok ? "Disconnected" : "Connection lost")
            appState?.onLiveSessionDisconnected()
        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's bootstrap/event-handling sync. 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's handling
    /// of UserEvent.UPDATED for the self user.
    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 joinVoice() {
        AVAudioApplication.requestRecordPermission { [weak self] granted in
            DispatchQueue.main.async {
                guard let self else { return }
                if granted {
                    let result = self.client.joinVoice()
                    if result != .ok {
                        self.addActivity("Failed to join voice: \(result.description)")
                    }
                } 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
        }

        let desc = StreamDescriptor(kind: .mic, deviceId: voiceState.currentDeviceId, label: "Mic",
                                    externalFeed: true)
        let (result, streamId) = client.startStream(desc)
        guard result == .ok else {
            addActivity("Failed to start mic: \(result.description)")
            return
        }
        voiceState.micActive = true
        voiceState.localStreamId = streamId
        EventFeedback.shared.play(.voiceOn)

        let channels = IOSAudioRouter.shared.captureChannels.channelCount
        if channels != 1 {
            client.setCaptureChannels(streamId: streamId, channels: channels)
        }
        IOSAudioEngine.shared.startMic(streamId: streamId, channels: channels)
    }

    func leaveVoice() {
        if voiceState.screenStreamId != 0 { stopScreenShare() }
        client.setPushToTalk(false)
        IOSAudioEngine.shared.stopMic()
        client.leaveVoice()
    }

    // MARK: - Reconnect restore

    /// Called by `AppState` after a reconnect's auth success to rejoin the prior channel and
    /// re-enable the prior voice/mic state. Drives the restore through the `.joinResult` event
    /// so we re-arm voice only AFTER the server processed the join — joining voice before the
    /// channel move would be rejected server-side. `micMuted`/`deafened` are the user's LOCAL
    /// mute/deafen state at the moment of the drop; the server resets those on a fresh auth, so
    /// we re-push them via `setMute` after the channel is restored.
    func requestRestore(channelId: UInt32, voiceSubscribed: Bool,
                        micMuted: Bool, deafened: Bool) {
        pendingRestore = RestoreRequest(channelId: channelId,
                                        voiceSubscribed: voiceSubscribed,
                                        micMuted: micMuted,
                                        deafened: deafened)
        didIssueRestoreJoin = false
        if channelId != 0 {
            // The server auto-placed us in the Lobby on auth; join our prior channel explicitly.
            // `vc_join_channel` is idempotent server-side (joining the channel you're already in
            // returns ok), so this is safe even if the prior channel was the Lobby.
            client.joinChannel(channelId)
            didIssueRestoreJoin = true
        } else {
            // No prior channel — go straight to the voice/mute restore. (voiceSubscribed with
            // channelId == 0 is contradictory; `completeRestore` further guards on
            // currentChannelId != 0 before subscribing to voice.)
            completeRestore()
        }
    }

    /// Finish the restore after the channel is in place (or there was no channel to restore):
    /// re-subscribe to voice if the user was transmitting, and re-apply the local mute/deafen
    /// state. Safe to call once per `pendingRestore`; clears it.
    private func completeRestore() {
        guard let r = pendingRestore else { return }
        if r.voiceSubscribed && currentChannelId != 0 {
            joinVoice()
        }
        setMute(r.micMuted, deafened: r.deafened)
        addActivity("Restored to channel \(currentChannelId)"
                    + (r.voiceSubscribed ? " with voice" : ""))
        pendingRestore = nil
    }

    // 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
        UserDefaults.standard.set(Int(mode.rawValue), forKey: DefaultsKey.inputMode)
    }

    func setVadThreshold(_ threshold: Float) {
        client.setVadThreshold(threshold)
        voiceState.vadThreshold = threshold
        UserDefaults.standard.set(threshold, forKey: DefaultsKey.vadThreshold)
    }

    func setInputGain(_ gain: Float) {
        client.setInputGain(gain)
        voiceState.inputGain = gain
        UserDefaults.standard.set(gain, forKey: DefaultsKey.inputGain)
    }

    func setInputNoiseReduction(_ on: Bool) {
        client.setInputNoiseReduction(on)
        voiceState.inputNoiseReduction = on
        UserDefaults.standard.set(on, forKey: DefaultsKey.inputNoiseReduction)
    }

    // MARK: - Persisted input settings

    private enum DefaultsKey {
        static let inputMode    = "voice.inputMode"
        static let vadThreshold = "voice.vadThreshold"
        static let inputGain    = "voice.inputGain"
        static let inputNoiseReduction = "voice.inputNoiseReduction"
    }

    /// Restore the saved input mode / VAD threshold / mic gain and push them into the core so a
    /// relaunch keeps the user's transmission settings instead of resetting to VAD defaults.
    private func loadAndApplyVoiceSettings() {
        let d = UserDefaults.standard
        if d.object(forKey: DefaultsKey.inputMode) != nil {
            let raw = UInt32(d.integer(forKey: DefaultsKey.inputMode))
            voiceState.inputMode = VoiceCatInputMode(rawValue: raw) ?? .voiceActivation
        }
        if d.object(forKey: DefaultsKey.vadThreshold) != nil {
            voiceState.vadThreshold = d.float(forKey: DefaultsKey.vadThreshold)
        }
        if d.object(forKey: DefaultsKey.inputGain) != nil {
            voiceState.inputGain = d.float(forKey: DefaultsKey.inputGain)
        }
        if d.object(forKey: DefaultsKey.inputNoiseReduction) != nil {
            voiceState.inputNoiseReduction = d.bool(forKey: DefaultsKey.inputNoiseReduction)
        }
        client.setInputMode(voiceState.inputMode)
        client.setVadThreshold(voiceState.vadThreshold)
        client.setInputGain(voiceState.inputGain)
        client.setInputNoiseReduction(voiceState.inputNoiseReduction)
    }

    private var pttEngaged = false
    func setPushToTalk(_ active: Bool) {
        client.setPushToTalk(active)
        // Play the PTT cue only on the press transition (the gesture fires repeatedly while held).
        if active && !pttEngaged { EventFeedback.shared.play(.ptt) }
        pttEngaged = 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)
    }
}