//
//  BPMDetector.swift
//  DashSmash
//
//  Tries hard to find a song's BPM so the player doesn't have to tap-tempo
//  every time. The cascade is, in order:
//
//   1. A per-source persistent cache (anything we've ever detected, or that
//      the player tapped manually and confirmed).
//   2. The MPMediaItem `beatsPerMinute` property — set by iTunes/Music.app
//      for well-tagged libraries.
//   3. Embedded MP3 / iTunes file metadata (`TBPM`, "iTunes BPM", etc.).
//   4. Audio analysis: read PCM from the middle of the song, build an
//      onset-novelty curve, autocorrelate, and pick the lag with the
//      strongest periodicity in the 60–180 BPM band.
//
//  Anything outside [60, 180] BPM is rejected as a likely false positive.
//  Tap-tempo remains the final fallback when none of these work.
//

import Foundation
import AVFoundation
import MediaPlayer

final class BPMDetector {
    static let shared = BPMDetector()

    enum DetectionSource: String {
        case cache
        case metadata
        case fileMetadata
        case analysis
    }

    struct Detection {
        let bpm: Double
        let source: DetectionSource
    }

    private let queue = DispatchQueue(label: "DashSmash.BPMDetector", qos: .userInitiated)

    /// Tries the full detection cascade for a track. Returns nil if every
    /// strategy fails or the candidate BPM is outside the usable range, in
    /// which case the caller should fall back to tap-tempo.
    func detect(for track: PlayableTrack) async -> Detection? {
        if let cached = BPMCache.shared.bpm(for: track.sourceID) {
            return Detection(bpm: cached, source: .cache)
        }
        switch track {
        case .library(let item):
            if let bpm = libraryMetadataBPM(item), inRange(bpm) {
                BPMCache.shared.setBPM(bpm, for: track.sourceID)
                return Detection(bpm: bpm, source: .metadata)
            }
            if let url = item.assetURL {
                if let bpm = await fileMetadataBPM(at: url), inRange(bpm) {
                    BPMCache.shared.setBPM(bpm, for: track.sourceID)
                    return Detection(bpm: bpm, source: .fileMetadata)
                }
                if let bpm = await analyzeBPM(at: url), inRange(bpm) {
                    BPMCache.shared.setBPM(bpm, for: track.sourceID)
                    return Detection(bpm: bpm, source: .analysis)
                }
            }
        case .appleMusicCatalog, .midi:
            // Catalog tracks have no local asset URL we can decode; MIDI is
            // already a self-tempoed format and shouldn't reach this path.
            break
        }
        return nil
    }

    /// Lets the tap-tempo path bypass the detector but still feed the cache,
    /// so the next time this song is picked we skip tap-tempo entirely.
    func remember(bpm: Double, for sourceID: String) {
        guard inRange(bpm) else { return }
        BPMCache.shared.setBPM(bpm, for: sourceID)
    }

    private func inRange(_ bpm: Double) -> Bool {
        bpm >= 50 && bpm <= 220
    }

    // MARK: - MPMediaItem metadata

    private func libraryMetadataBPM(_ item: MPMediaItem) -> Double? {
        if let number = item.value(forProperty: MPMediaItemPropertyBeatsPerMinute) as? NSNumber {
            let value = number.doubleValue
            if value >= 30 { return value }
        }
        return nil
    }

    // MARK: - AVAsset file metadata

    private func fileMetadataBPM(at url: URL) async -> Double? {
        let asset = AVURLAsset(url: url)
        let formats: [AVMetadataFormat] = [.id3Metadata, .iTunesMetadata, .quickTimeMetadata, .quickTimeUserData]
        for format in formats {
            guard let items = try? await asset.loadMetadata(for: format) else { continue }
            for item in items {
                if let bpm = await extractBPM(from: item) { return bpm }
            }
        }
        // Common-key fallback
        if let common = try? await asset.load(.commonMetadata) {
            for item in common {
                if let bpm = await extractBPM(from: item) { return bpm }
            }
        }
        return nil
    }

    private func extractBPM(from item: AVMetadataItem) async -> Double? {
        let keyString = (item.key as? String) ?? ""
        let identifier = item.identifier?.rawValue ?? ""
        let hits = ["TBPM", "tbpm", "BPM", "bpm", "beatsPerMinute", "tmpo"]
        let matches = hits.contains(where: { keyString.localizedCaseInsensitiveContains($0) || identifier.localizedCaseInsensitiveContains($0) })
        guard matches else { return nil }
        if let number = try? await item.load(.numberValue) {
            return number.doubleValue
        }
        if let str = try? await item.load(.stringValue), let v = Double(str) {
            return v
        }
        return nil
    }

    // MARK: - Audio analysis

    /// Reads ~30 seconds of mono PCM from roughly the middle of the song and
    /// runs an onset-novelty autocorrelation to find the dominant tempo.
    /// Returns nil if the track can't be decoded, is too short, or the
    /// autocorrelation peak is too weak to trust.
    private func analyzeBPM(at url: URL) async -> Double? {
        await withCheckedContinuation { continuation in
            queue.async {
                let result = self.analyzeBPMSync(at: url)
                continuation.resume(returning: result)
            }
        }
    }

    private func analyzeBPMSync(at url: URL) -> Double? {
        let asset = AVURLAsset(url: url)
        let durationCM = asset.duration
        let durationSeconds = CMTimeGetSeconds(durationCM)
        guard durationSeconds > 20 else { return nil }

        // Anchor analysis in a region likely to have a stable beat — past
        // the intro, before any outro. Falls back to the start for short
        // tracks.
        let analysisLength: Double = min(30, max(15, durationSeconds * 0.4))
        let middleStart = max(0, min(durationSeconds - analysisLength, durationSeconds * 0.33))

        guard let audioTrack = asset.tracks(withMediaType: .audio).first else { return nil }
        guard let reader = try? AVAssetReader(asset: asset) else { return nil }

        let sampleRate: Double = 22050
        let settings: [String: Any] = [
            AVFormatIDKey: kAudioFormatLinearPCM,
            AVSampleRateKey: sampleRate,
            AVNumberOfChannelsKey: 1,
            AVLinearPCMBitDepthKey: 16,
            AVLinearPCMIsBigEndianKey: false,
            AVLinearPCMIsFloatKey: false,
            AVLinearPCMIsNonInterleaved: false
        ]
        let output = AVAssetReaderTrackOutput(track: audioTrack, outputSettings: settings)
        output.alwaysCopiesSampleData = false
        guard reader.canAdd(output) else { return nil }
        reader.add(output)
        reader.timeRange = CMTimeRange(
            start: CMTime(seconds: middleStart, preferredTimescale: 44100),
            duration: CMTime(seconds: analysisLength, preferredTimescale: 44100)
        )
        guard reader.startReading() else { return nil }

        var samples: [Int16] = []
        samples.reserveCapacity(Int(analysisLength * sampleRate))
        while let buffer = output.copyNextSampleBuffer() {
            if let block = CMSampleBufferGetDataBuffer(buffer) {
                var length = 0
                var pointer: UnsafeMutablePointer<Int8>?
                CMBlockBufferGetDataPointer(block, atOffset: 0, lengthAtOffsetOut: nil,
                                            totalLengthOut: &length, dataPointerOut: &pointer)
                if let pointer {
                    let count = length / MemoryLayout<Int16>.size
                    pointer.withMemoryRebound(to: Int16.self, capacity: count) { i16 in
                        samples.append(contentsOf: UnsafeBufferPointer(start: i16, count: count))
                    }
                }
            }
            CMSampleBufferInvalidate(buffer)
        }
        reader.cancelReading()

        return computeBPM(samples: samples, sampleRate: sampleRate)
    }

    private func computeBPM(samples: [Int16], sampleRate: Double) -> Double? {
        let hopSize = 512
        let frameCount = samples.count / hopSize
        guard frameCount > 64 else { return nil }

        // Frame-wise RMS envelope.
        var envelope = [Float](repeating: 0, count: frameCount)
        for f in 0..<frameCount {
            var sum: Float = 0
            let base = f * hopSize
            for i in 0..<hopSize {
                let s = Float(samples[base + i]) / 32768.0
                sum += s * s
            }
            envelope[f] = sqrt(sum / Float(hopSize))
        }

        // Onset novelty = half-wave rectified first difference.
        var novelty = [Float](repeating: 0, count: frameCount)
        for i in 1..<frameCount {
            let d = envelope[i] - envelope[i - 1]
            novelty[i] = max(0, d)
        }

        // Zero-mean the novelty signal so autocorrelation peaks reflect
        // periodicity rather than DC offset.
        let mean = novelty.reduce(0, +) / Float(frameCount)
        for i in 0..<frameCount { novelty[i] -= mean }

        let frameRate = Float(sampleRate) / Float(hopSize)
        let lagMin = max(1, Int((60.0 * Double(frameRate) / 200.0).rounded()))
        let lagMax = min(frameCount - 1, Int((60.0 * Double(frameRate) / 60.0).rounded()))
        guard lagMax > lagMin else { return nil }

        var bestLag = 0
        var bestScore: Float = -.infinity
        for lag in lagMin...lagMax {
            var sum: Float = 0
            let limit = frameCount - lag
            for i in 0..<limit {
                sum += novelty[i] * novelty[i + lag]
            }
            if sum > bestScore {
                bestScore = sum
                bestLag = lag
            }
        }
        guard bestLag > 0 else { return nil }

        // Reject if the autocorrelation peak is barely above zero — that's
        // a signal we couldn't find a stable beat.
        let normalizer = max(novelty.map { $0 * $0 }.reduce(0, +), 1e-6)
        let confidence = bestScore / normalizer
        guard confidence > 0.02 else { return nil }

        let rawBPM = 60.0 * Double(frameRate) / Double(bestLag)
        // Snap into a comfortable range: if we detected something in the
        // 60–80 band but the actual energy at half-lag is similar, prefer
        // the doubled tempo (most modern songs sit in 95–160).
        return Self.preferredOctave(rawBPM, novelty: novelty, bestLag: bestLag)
    }

    /// If we picked a slow tempo whose double is also strongly periodic,
    /// prefer the double — most pop/dance songs feel right in the 95–160 band.
    private static func preferredOctave(_ bpm: Double, novelty: [Float], bestLag: Int) -> Double {
        guard bpm < 90 else { return bpm }
        let halfLag = bestLag / 2
        guard halfLag > 1, halfLag < novelty.count - 1 else { return bpm }
        var sum: Float = 0
        let limit = novelty.count - halfLag
        for i in 0..<limit { sum += novelty[i] * novelty[i + halfLag] }
        // If half-lag correlation is at least 70% as strong as the chosen
        // lag, the song is probably a double-tempo at this measurement.
        var fullSum: Float = 0
        let fullLimit = novelty.count - bestLag
        for i in 0..<fullLimit { fullSum += novelty[i] * novelty[i + bestLag] }
        if fullSum > 0 && sum / fullSum > 0.7 {
            return bpm * 2
        }
        return bpm
    }
}

// MARK: - Persistent cache

/// Simple UserDefaults-backed map of `sourceID → bpm`. Keys are the same
/// stable identifiers `PlayableTrack.sourceID` uses, so library, catalog,
/// and MIDI tracks can coexist without colliding.
final class BPMCache {
    static let shared = BPMCache()

    private let key = "DashSmash.bpmCache.v1"
    private let defaults: UserDefaults

    private init(defaults: UserDefaults = .standard) {
        self.defaults = defaults
    }

    func bpm(for sourceID: String) -> Double? {
        let map = defaults.dictionary(forKey: key) as? [String: Double] ?? [:]
        return map[sourceID]
    }

    func setBPM(_ bpm: Double, for sourceID: String) {
        var map = defaults.dictionary(forKey: key) as? [String: Double] ?? [:]
        map[sourceID] = bpm
        defaults.set(map, forKey: key)
    }
}