ios-celestia-trace/CelestiaTrace/Services/Audio/SilenceDetector.swift

import Foundation
import AVFoundation

/// Represents a detected silent interval within an audio recording.
public struct SilenceRange: Codable, Hashable, Sendable {
    public let start: TimeInterval // in seconds
    public let end: TimeInterval   // in seconds
    
    public var duration: TimeInterval {
        end - start
    }
    
    public init(start: TimeInterval, end: TimeInterval) {
        self.start = start
        self.end = end
    }
}

/// Playback analysis generated in one pass over the decoded audio.
public struct AudioAnalysisResult: Sendable {
    public let waveformSamples: [Float]
    public let silentRanges: [SilenceRange]

    public init(waveformSamples: [Float], silentRanges: [SilenceRange]) {
        self.waveformSamples = waveformSamples
        self.silentRanges = silentRanges
    }
}

/// One progressively decoded section of a recording.
public struct AudioAnalysisProgress: Sendable {
    public let totalSampleCount: Int
    public let sampleOffset: Int
    public let waveformSamples: [Float]
    public let silentRanges: [SilenceRange]
    public let isComplete: Bool
}

private struct AudioAnalysisCacheEntry: Codable {
    let version: Int
    let fileSize: Int64
    let modificationTime: TimeInterval
    let waveformSamples: [Float]
    let silentRanges: [SilenceRange]
}

/// A cancellable, priority-aware analysis run. The encoded recording stays as
/// one file; only decoding is divided into logical 60-second sections.
public final class AudioAnalysisSession: @unchecked Sendable {
    private static let cacheVersion = 1
    private static let sampleInterval: TimeInterval = 0.05
    private static let segmentDuration: TimeInterval = 60

    private let audioURL: URL
    private let silenceThresholdDB: Float
    private let minimumSilenceDuration: TimeInterval
    private let stateLock = NSLock()
    private var cancelled = false
    private var prioritizedTime: TimeInterval?

    public init(
        audioURL: URL,
        silenceThresholdDB: Float = -40,
        minimumSilenceDuration: TimeInterval = 2
    ) {
        self.audioURL = audioURL
        self.silenceThresholdDB = silenceThresholdDB
        self.minimumSilenceDuration = minimumSilenceDuration
    }

    public func cancel() {
        stateLock.lock()
        cancelled = true
        stateLock.unlock()
    }

    public func prioritize(time: TimeInterval) {
        stateLock.lock()
        prioritizedTime = max(0, time)
        stateLock.unlock()
    }

    public func run(
        onUpdate: @escaping @Sendable (AudioAnalysisProgress) async -> Void
    ) async {
        await Task.detached(priority: .utility) { [self] in
            await runDetached(onUpdate: onUpdate)
        }.value
    }

    private func runDetached(
        onUpdate: @escaping @Sendable (AudioAnalysisProgress) async -> Void
    ) async {
        guard let validURL = AudioPathHelper.resolveURL(for: audioURL.path),
              let identity = Self.fileIdentity(for: validURL) else {
            return
        }

        if usesDefaultCacheSettings,
           let cached = Self.loadCache(for: validURL, identity: identity) {
            await onUpdate(AudioAnalysisProgress(
                totalSampleCount: cached.waveformSamples.count,
                sampleOffset: 0,
                waveformSamples: cached.waveformSamples,
                silentRanges: cached.silentRanges,
                isComplete: true
            ))
            return
        }

        guard let audioFile = try? AVAudioFile(forReading: validURL) else { return }
        let format = audioFile.processingFormat
        let sampleRate = format.sampleRate
        guard sampleRate > 0 else { return }

        let subChunkFrames = max(1, Int64(sampleRate * Self.sampleInterval))
        let segmentFrames = max(subChunkFrames, Int64(sampleRate * Self.segmentDuration))
        let totalFrames = audioFile.length
        let totalSampleCount = Int((totalFrames + subChunkFrames - 1) / subChunkFrames)
        let totalSegmentCount = max(1, Int((totalFrames + segmentFrames - 1) / segmentFrames))

        var completedSegments = Set<Int>()
        var completeWaveform = Array(repeating: Float(0), count: totalSampleCount)
        var rawSilentRanges: [SilenceRange] = []

        while completedSegments.count < totalSegmentCount {
            if isCancelled || Task.isCancelled { return }
            let segmentIndex = nextSegmentIndex(
                totalSegmentCount: totalSegmentCount,
                completedSegments: completedSegments
            )
            guard let segmentIndex else { break }

            let startFrame = Int64(segmentIndex) * segmentFrames
            let endFrame = min(startFrame + segmentFrames, totalFrames)
            guard let segment = Self.analyzeSegment(
                audioFile: audioFile,
                format: format,
                sampleRate: sampleRate,
                startFrame: startFrame,
                endFrame: endFrame,
                subChunkFrames: subChunkFrames,
                silenceThresholdDB: silenceThresholdDB
            ) else {
                return
            }

            completedSegments.insert(segmentIndex)
            let sampleOffset = Int(startFrame / subChunkFrames)
            let upperBound = min(sampleOffset + segment.waveformSamples.count, completeWaveform.count)
            if sampleOffset < upperBound {
                completeWaveform.replaceSubrange(
                    sampleOffset..<upperBound,
                    with: segment.waveformSamples.prefix(upperBound - sampleOffset)
                )
            }
            rawSilentRanges.append(contentsOf: segment.rawSilentRanges)
            let mergedSilence = SilenceDetector.mergeSilentRanges(
                rawSilentRanges,
                minimumDuration: minimumSilenceDuration
            )
            let isComplete = completedSegments.count == totalSegmentCount

            await onUpdate(AudioAnalysisProgress(
                totalSampleCount: totalSampleCount,
                sampleOffset: sampleOffset,
                waveformSamples: segment.waveformSamples,
                silentRanges: mergedSilence,
                isComplete: isComplete
            ))
            await Task.yield()

            if isComplete, usesDefaultCacheSettings {
                Self.saveCache(
                    AudioAnalysisCacheEntry(
                        version: Self.cacheVersion,
                        fileSize: identity.fileSize,
                        modificationTime: identity.modificationTime,
                        waveformSamples: completeWaveform,
                        silentRanges: mergedSilence
                    ),
                    for: validURL
                )
            }
        }
    }

    private var isCancelled: Bool {
        stateLock.lock()
        defer { stateLock.unlock() }
        return cancelled
    }

    private var usesDefaultCacheSettings: Bool {
        silenceThresholdDB == -40 && minimumSilenceDuration == 2
    }

    private func nextSegmentIndex(
        totalSegmentCount: Int,
        completedSegments: Set<Int>
    ) -> Int? {
        stateLock.lock()
        let preferredTime = prioritizedTime
        prioritizedTime = nil
        stateLock.unlock()

        if let preferredTime {
            let preferredIndex = min(
                max(Int(preferredTime / Self.segmentDuration), 0),
                totalSegmentCount - 1
            )
            if !completedSegments.contains(preferredIndex) {
                return preferredIndex
            }
            // After the requested section, favor its immediate neighbors.
            for distance in 1..<totalSegmentCount {
                let forward = preferredIndex + distance
                if forward < totalSegmentCount, !completedSegments.contains(forward) {
                    return forward
                }
                let backward = preferredIndex - distance
                if backward >= 0, !completedSegments.contains(backward) {
                    return backward
                }
            }
        }
        return (0..<totalSegmentCount).first { !completedSegments.contains($0) }
    }

    private struct SegmentResult {
        let waveformSamples: [Float]
        let rawSilentRanges: [SilenceRange]
    }

    private static func analyzeSegment(
        audioFile: AVAudioFile,
        format: AVAudioFormat,
        sampleRate: Double,
        startFrame: Int64,
        endFrame: Int64,
        subChunkFrames: Int64,
        silenceThresholdDB: Float
    ) -> SegmentResult? {
        let oneSecondFrames = max(1, AVAudioFrameCount(sampleRate))
        guard let buffer = AVAudioPCMBuffer(
            pcmFormat: format,
            frameCapacity: oneSecondFrames
        ) else {
            return nil
        }

        audioFile.framePosition = startFrame
        var waveformSamples: [Float] = []
        waveformSamples.reserveCapacity(Int((endFrame - startFrame + subChunkFrames - 1) / subChunkFrames))
        var rawSilentRanges: [SilenceRange] = []
        var silenceStart: TimeInterval?

        do {
            while audioFile.framePosition < endFrame {
                if Task.isCancelled { return nil }
                let remainingFrames = endFrame - audioFile.framePosition
                let framesToRead = AVAudioFrameCount(
                    min(Int64(oneSecondFrames), remainingFrames)
                )
                guard framesToRead > 0 else { break }
                try audioFile.read(into: buffer, frameCount: framesToRead)
                guard let channelData = buffer.floatChannelData?[0] else { continue }

                let bufferStartFrame = audioFile.framePosition - Int64(buffer.frameLength)
                var offset = 0
                while offset < Int(buffer.frameLength) {
                    let sampleCount = min(Int(subChunkFrames), Int(buffer.frameLength) - offset)
                    guard sampleCount > 0 else { break }
                    var sum: Float = 0
                    for index in 0..<sampleCount {
                        let sample = channelData[offset + index]
                        sum += sample * sample
                    }

                    let rms = sqrt(sum / Float(sampleCount))
                    let decibels = rms > 0 ? 20 * log10(rms) : -100
                    waveformSamples.append(
                        AudioLevelNormalizer.normalizedLevel(decibels: decibels)
                    )

                    let absoluteFrame = bufferStartFrame + Int64(offset)
                    let time = Double(absoluteFrame) / sampleRate
                    if decibels < silenceThresholdDB {
                        if silenceStart == nil {
                            silenceStart = time
                        }
                    } else if let start = silenceStart {
                        rawSilentRanges.append(SilenceRange(start: start, end: time))
                        silenceStart = nil
                    }
                    offset += sampleCount
                }
            }
        } catch {
            print("[SilenceDetector] Error reading audio segment: \(error.localizedDescription)")
            return nil
        }

        if let start = silenceStart {
            rawSilentRanges.append(
                SilenceRange(start: start, end: Double(endFrame) / sampleRate)
            )
        }
        return SegmentResult(
            waveformSamples: waveformSamples,
            rawSilentRanges: rawSilentRanges
        )
    }

    private struct FileIdentity {
        let fileSize: Int64
        let modificationTime: TimeInterval
    }

    private static func fileIdentity(for url: URL) -> FileIdentity? {
        guard let values = try? url.resourceValues(
            forKeys: [.fileSizeKey, .contentModificationDateKey]
        ) else {
            return nil
        }
        return FileIdentity(
            fileSize: Int64(values.fileSize ?? 0),
            modificationTime: values.contentModificationDate?.timeIntervalSince1970 ?? 0
        )
    }

    private static func cacheURL(for audioURL: URL) -> URL? {
        guard let cacheRoot = FileManager.default.urls(
            for: .cachesDirectory,
            in: .userDomainMask
        ).first else {
            return nil
        }
        let directory = cacheRoot
            .appendingPathComponent("CelestiaTrace", isDirectory: true)
            .appendingPathComponent("AudioAnalysis", isDirectory: true)
        try? FileManager.default.createDirectory(
            at: directory,
            withIntermediateDirectories: true
        )

        var hash: UInt64 = 1_469_598_103_934_665_603
        // Persist across app-container path changes by keying on the stable
        // recording filename; size and modification time validate the content.
        for byte in audioURL.lastPathComponent.utf8 {
            hash ^= UInt64(byte)
            hash &*= 1_099_511_628_211
        }
        return directory.appendingPathComponent(String(hash, radix: 16) + ".plist")
    }

    private static func loadCache(
        for audioURL: URL,
        identity: FileIdentity
    ) -> AudioAnalysisCacheEntry? {
        guard let cacheURL = cacheURL(for: audioURL),
              let data = try? Data(contentsOf: cacheURL),
              let entry = try? PropertyListDecoder().decode(
                AudioAnalysisCacheEntry.self,
                from: data
              ),
              entry.version == cacheVersion,
              entry.fileSize == identity.fileSize,
              abs(entry.modificationTime - identity.modificationTime) < 0.001 else {
            return nil
        }
        return entry
    }

    private static func saveCache(
        _ entry: AudioAnalysisCacheEntry,
        for audioURL: URL
    ) {
        guard let cacheURL = cacheURL(for: audioURL) else { return }
        let encoder = PropertyListEncoder()
        encoder.outputFormat = .binary
        guard let data = try? encoder.encode(entry) else { return }
        try? data.write(to: cacheURL, options: .atomic)
    }
}

private actor AudioAnalysisCollector {
    private var waveformSamples: [Float] = []
    private var silentRanges: [SilenceRange] = []

    func apply(_ update: AudioAnalysisProgress) {
        if waveformSamples.count != update.totalSampleCount {
            waveformSamples = Array(repeating: 0, count: update.totalSampleCount)
        }
        let lowerBound = min(max(update.sampleOffset, 0), waveformSamples.count)
        let upperBound = min(lowerBound + update.waveformSamples.count, waveformSamples.count)
        if lowerBound < upperBound {
            waveformSamples.replaceSubrange(
                lowerBound..<upperBound,
                with: update.waveformSamples.prefix(upperBound - lowerBound)
            )
        }
        silentRanges = update.silentRanges
    }

    func result() -> AudioAnalysisResult {
        AudioAnalysisResult(
            waveformSamples: waveformSamples,
            silentRanges: silentRanges
        )
    }
}

public final class SilenceDetector: Sendable {
    /// Detects silent ranges in the specified audio file.
    /// - Parameters:
    ///   - audioURL: Local file URL of the audio.
    ///   - thresholdDB: Threshold in decibels (e.g. -40.0 dB). Sounds below this are considered silent.
    ///   - minDuration: Minimum consecutive duration in seconds to qualify as a silent segment.
    /// - Returns: An array of detected SilenceRange objects.
    public static func detectSilence(
        in audioURL: URL,
        thresholdDB: Float = -40.0,
        minDuration: TimeInterval = 2.0
    ) async -> [SilenceRange] {
        await analyze(
            audioURL,
            silenceThresholdDB: thresholdDB,
            minimumSilenceDuration: minDuration
        ).silentRanges
    }

    /// Extracts waveform levels and silence ranges together to avoid decoding twice.
    /// Waveform samples use the same 50ms cadence and dB normalization as live recording.
    public static func analyze(
        _ audioURL: URL,
        silenceThresholdDB: Float = -40.0,
        minimumSilenceDuration: TimeInterval = 2.0
    ) async -> AudioAnalysisResult {
        let collector = AudioAnalysisCollector()
        let session = AudioAnalysisSession(
            audioURL: audioURL,
            silenceThresholdDB: silenceThresholdDB,
            minimumSilenceDuration: minimumSilenceDuration
        )
        await withTaskCancellationHandler {
            await session.run { update in
                await collector.apply(update)
            }
        } onCancel: {
            session.cancel()
        }
        return await collector.result()
    }

    /// Warms the derived cache without retaining waveform data in a view model.
    public static func warmCache(for audioURL: URL) async {
        let session = AudioAnalysisSession(audioURL: audioURL)
        await withTaskCancellationHandler {
            await session.run { _ in }
        } onCancel: {
            session.cancel()
        }
    }

    /// Merges silence that crosses a logical segment boundary, then applies the
    /// product's two-second minimum. Kept internal for focused logic tests.
    static func mergeSilentRanges(
        _ ranges: [SilenceRange],
        minimumDuration: TimeInterval = 2
    ) -> [SilenceRange] {
        let sorted = ranges.sorted {
            if $0.start == $1.start {
                return $0.end < $1.end
            }
            return $0.start < $1.start
        }
        guard var current = sorted.first else { return [] }
        var merged: [SilenceRange] = []
        for next in sorted.dropFirst() {
            if next.start <= current.end + 0.001 {
                current = SilenceRange(
                    start: current.start,
                    end: max(current.end, next.end)
                )
            } else {
                if current.duration >= minimumDuration {
                    merged.append(current)
                }
                current = next
            }
        }
        if current.duration >= minimumDuration {
            merged.append(current)
        }
        return merged
    }
}