import AVFoundation
import Foundation

@inline(__always) func monoNanos() -> Int64 { Int64(DispatchTime.now().uptimeNanoseconds) }

/// One track's stems, decoding in sync. Built off the main thread, then handed to the engine.
final class TrackSession: @unchecked Sendable {
    let token: Int64
    let names: [String]
    let decoders: [StemDecoder]
    var sampleRate: Int { decoders.first?.sampleRate ?? 0 }
    var durationFrames: Int64 { (decoders.map(\.durationUs).max() ?? 0) * Int64(sampleRate) / 1_000_000 }
    fileprivate var gen = 0
    fileprivate var framePos: Int64 = 0
    fileprivate var target: [Float]
    fileprivate var current: [Float]
    fileprivate var mixVersion = -1

    init(token: Int64, names: [String], decoders: [StemDecoder]) {
        self.token = token
        self.names = names
        self.decoders = decoders
        target = Array(repeating: 1, count: names.count)
        current = Array(repeating: 1, count: names.count)
    }

    func error() -> Error? { decoders.lazy.compactMap(\.error).first }

    fileprivate func ready(_ need: Int) -> Bool {
        decoders.allSatisfy { d in
            d.error != nil || (d.readyGen == gen || d.readyGen == Int.max) && (d.ring.eos || d.ring.availableFrames >= need)
        }
    }

    fileprivate func finished() -> Bool {
        decoders.allSatisfy { d in
            d.error != nil || (d.readyGen == gen || d.readyGen == Int.max) && d.ring.eos && d.ring.availableFrames == 0
        }
    }

    /// Re-targets decoders; only valid before the session is handed to the engine.
    func startAt(_ frame: Int64) {
        gen += 1
        framePos = frame
        decoders.forEach { $0.seek(gen: gen, frame: frame) }
    }

    func release() { decoders.forEach { $0.release() } }
}

/// Per-stem RMS of short blocks as they're mixed, stamped with their track frame, so meters can
/// show what's coming out of the speaker now rather than what was just mixed.
final class MeterHistory: @unchecked Sendable {
    static let block = 256
    static let maxStems = 8
    /// ~5 s at 48 kHz: well past the output buffer.
    private static let cap = 1024
    private static let staleFrames: Int64 = 8192

    private var tokens = [Int64](repeating: 0, count: cap)
    private var starts = [Int64](repeating: 0, count: cap)
    private var values = [Float](repeating: 0, count: cap * maxStems)
    private var head = 0
    private var count = 0
    private let lock = UnfairLock()

    func clear() { lock.with { count = 0 } }

    func add(_ token: Int64, _ start: Int64, _ rms: UnsafePointer<Float>, _ stems: Int) {
        lock.with {
            tokens[head] = token
            starts[head] = start
            let base = head * Self.maxStems
            for k in 0..<Self.maxStems { values[base + k] = k < stems ? rms[k] : 0 }
            head = (head + 1) % Self.cap
            count = min(count + 1, Self.cap)
        }
    }

    /// Fills [out] with the levels of [token]'s block playing at [frame]; false if there's none.
    func read(_ token: Int64, _ frame: Int64, _ out: inout [Float]) -> Bool {
        lock.with {
            guard count > 0 else { return false }
            for k in 1...count {
                let i = (head - k + Self.cap) % Self.cap
                if tokens[i] != token || starts[i] > frame { continue }
                if frame - starts[i] > Self.staleFrames { return false }
                for j in 0..<min(out.count, Self.maxStems) { out[j] = values[i * Self.maxStems + j] }
                return true
            }
            return false
        }
    }
}

/// The far end of the output ring: the render callback takes mixed frames from it and counts
/// what it took, so the engine knows what has reached the speaker.
private final class OutputTap: @unchecked Sendable {
    let ring: FloatRing
    let rate: Int
    private let lock = UnfairLock()
    private var consumed: Int64 = 0
    private var lastStart: Int64 = 0
    private var lastNanos: Int64 = 0
    private var runningFlag = false
    private var latencyNanos: Int64 = 0
    private let scratch: UnsafeMutablePointer<Float>
    private static let scratchFrames = 4096

    init(rate: Int, frames: Int) {
        self.rate = rate
        ring = FloatRing(capacityFrames: frames)
        scratch = .allocate(capacity: Self.scratchFrames * 2)
    }
    deinit { scratch.deallocate() }

    var running: Bool {
        get { lock.with { runningFlag } }
        set { lock.with { runningFlag = newValue } }
    }

    func setLatency(_ seconds: Double) { lock.with { latencyNanos = Int64(seconds * 1e9) } }

    func flush() {
        lock.with {
            ring.clear()
            consumed = 0
            lastStart = 0
            lastNanos = 0
        }
    }

    /// Frames heard so far (since the last flush), never more than [written].
    func presented(_ now: Int64, _ written: Int64) -> Int64 {
        lock.with {
            if lastNanos == 0 { return 0 }
            let elapsed = runningFlag ? max(0, now - lastNanos) : 0
            let est = lastStart + elapsed * Int64(rate) / 1_000_000_000 - latencyNanos * Int64(rate) / 1_000_000_000
            return min(max(0, min(est, consumed)), written)
        }
    }

    func render(_ silence: UnsafeMutablePointer<ObjCBool>, _ ts: UnsafePointer<AudioTimeStamp>, _ frames: AVAudioFrameCount, _ abl: UnsafeMutablePointer<AudioBufferList>) -> OSStatus {
        let list = UnsafeMutableAudioBufferListPointer(abl)
        let n = Int(frames)
        guard list.count >= 2, let l = list[0].mData?.assumingMemoryBound(to: Float.self),
              let r = list[1].mData?.assumingMemoryBound(to: Float.self) else { return noErr }
        lock.lock()
        if !runningFlag {
            lock.unlock()
            l.update(repeating: 0, count: n)
            r.update(repeating: 0, count: n)
            silence.pointee = true
            return noErr
        }
        let host = ts.pointee.mFlags.contains(.hostTimeValid) ? Int64(AVAudioTime.seconds(forHostTime: ts.pointee.mHostTime) * 1e9) : monoNanos()
        lastStart = consumed
        lastNanos = host
        var done = 0
        while done < n {
            let want = min(n - done, Self.scratchFrames)
            let got = ring.read(scratch, want)
            for i in 0..<got {
                l[done + i] = scratch[i * 2]
                r[done + i] = scratch[i * 2 + 1]
            }
            consumed += Int64(got)
            done += got
            if got < want { break }
        }
        lock.unlock()
        if done < n {
            (l + done).update(repeating: 0, count: n - done)
            (r + done).update(repeating: 0, count: n - done)
        }
        return noErr
    }
}

protocol StemEngineListener: AnyObject {
    func onEngineState()
    /// [token]'s audio is now what's coming out of the speaker.
    func onTransition(_ token: Int64)
    func onEnded(_ token: Int64)
    func onError(_ token: Int64, _ message: String)
}

/// Mixes a [TrackSession]'s stems into one output: per-stem gain (ramped), EQ with anti-clip
/// preamp, master gain, then a peak limiter. All output calls happen on the engine thread; the
/// public methods just queue commands.
final class StemEngine: @unchecked Sendable {
    static let chunk = 1024
    /// Per stem, decoded ahead of the mixer.
    static let ringSeconds = 3
    /// Frames each stem must have before (re)starting output.
    static let prefill = 8192
    static let bufferMs = 120

    /// Where playback is; position = frame + (now - nanos) while [running].
    struct Snapshot {
        var token: Int64
        var frame: Int64
        var rate: Int
        var nanos: Int64
        var running: Bool
        func positionMs(_ now: Int64 = monoNanos()) -> Int64 {
            if rate <= 0 { return 0 }
            let f = running ? frame + (now - nanos) * Int64(rate) / 1_000_000_000 : frame
            return f * 1000 / Int64(rate)
        }
    }

    private struct Segment { let writtenStart: Int64; let token: Int64; let frameStart: Int64 }

    private weak var listener: StemEngineListener?
    private let cmdLock = NSCondition()
    private var commands: [() -> Void] = []
    private var released = false

    private var session: TrackSession?
    private var next: TrackSession?
    private var wantPlay = false
    private let av = AVAudioEngine()
    private var node: AVAudioSourceNode?
    private var tap: OutputTap?
    private var trackRate = 0
    private var written: Int64 = 0
    private var segments: [Segment] = []
    private var drainUntil: Int64 = -1
    private var drainDeadline: Int64 = 0
    private var drainToken: Int64 = -1
    private var waitingSince: Int64 = 0
    private var errorReported: Int64 = -1
    private let eq = Equalizer()
    private let limiter = Limiter()
    private var masterCur: Float = 1
    private let out = UnsafeMutablePointer<Float>.allocate(capacity: chunk * 2)
    private let tmp = UnsafeMutablePointer<Float>.allocate(capacity: chunk * 2)
    private let meters = MeterHistory()
    private var meterSums = [Float](repeating: 0, count: chunk / MeterHistory.block * MeterHistory.maxStems)

    private let shared = UnfairLock()
    private var _eqParams: EqParams?
    private var _masterGain: Float = 1
    private var _mix: [String: Float] = [:]
    private var _mixVersion = 0
    private var _snapshot = Snapshot(token: -1, frame: 0, rate: 44100, nanos: 0, running: false)
    private var _playing = false
    private var _buffering = false

    var eqParams: EqParams? {
        get { shared.with { _eqParams } }
        set { shared.with { _eqParams = newValue } }
    }
    var masterGain: Float {
        get { shared.with { _masterGain } }
        set { shared.with { _masterGain = newValue } }
    }
    private(set) var snapshot: Snapshot {
        get { shared.with { _snapshot } }
        set { shared.with { _snapshot = newValue } }
    }
    var playing: Bool { shared.with { _playing } }
    var buffering: Bool { shared.with { _buffering } }

    init(listener: StemEngineListener) {
        self.listener = listener
        out.initialize(repeating: 0, count: Self.chunk * 2)
        tmp.initialize(repeating: 0, count: Self.chunk * 2)
        NotificationCenter.default.addObserver(forName: .AVAudioEngineConfigurationChange, object: av, queue: nil) { [weak self] _ in
            // The route or hardware format changed and the engine stopped; start it again.
            self?.post { self?.restartOutput() }
        }
        #if os(iOS)
        NotificationCenter.default.addObserver(forName: AVAudioSession.routeChangeNotification, object: nil, queue: nil) { [weak self] _ in
            self?.post { self?.tap?.setLatency(Self.outputLatency()) }
        }
        #endif
        let t = Thread { [weak self] in self?.loop() }
        t.name = "stem-engine"
        t.qualityOfService = .userInteractive
        t.start()
    }

    // ---- commands (any thread) --------------------------------------------------------

    private func post(_ cmd: @escaping () -> Void) {
        cmdLock.lock()
        commands.append(cmd)
        cmdLock.broadcast()
        cmdLock.unlock()
    }

    /// Replaces whatever is playing with [s], which must already be prepared.
    func load(_ s: TrackSession, play: Bool) {
        post { [self] in
            dropAll()
            session = s
            wantPlay = play
            resetOutput(s)
            // Starts out buffering so the step that begins playback always reports a change,
            // even when the last track was already playing.
            setFlags(false, play)
            publish()
        }
    }

    /// The track to continue into gaplessly when the current one ends (or nil).
    func queueNext(_ s: TrackSession?) {
        post { [self] in
            if next !== s { next?.release() }
            next = s
        }
    }

    func setPlaying(_ play: Bool) {
        post { [self] in
            wantPlay = play
            if !play { pauseOutput() }
            publish()
        }
    }

    func seekTo(_ ms: Int64) {
        post { [self] in
            guard let s = session else { return }
            let frame = min(max(ms * Int64(s.sampleRate) / 1000, 0), max(0, s.durationFrames - 1))
            s.gen += 1
            s.framePos = frame
            s.decoders.forEach { $0.seek(gen: s.gen, frame: frame) }
            resetOutput(s)
            publish()
        }
    }

    func stop() {
        post { [self] in
            dropAll()
            wantPlay = false
            publish()
        }
    }

    func setMix(_ gains: [String: Float]) {
        shared.with {
            _mix = gains
            _mixVersion += 1
        }
    }

    func release() {
        cmdLock.lock()
        released = true
        cmdLock.broadcast()
        cmdLock.unlock()
    }

    /// The meters' levels for what's audible right now, into [out] (one per stem). False when
    /// nothing is playing, so the meters can fall.
    func levelsNow(_ out: inout [Float]) -> Bool {
        let snap = snapshot
        if !snap.running || snap.rate <= 0 { return false }
        let frame = snap.frame + (monoNanos() - snap.nanos) * Int64(snap.rate) / 1_000_000_000
        return meters.read(snap.token, frame, &out)
    }

    // ---- engine thread ----------------------------------------------------------------

    private var isReleased: Bool { cmdLock.withLock { released } }

    private func dropAll() {
        session?.release()
        next?.release()
        session = nil
        next = nil
        drainUntil = -1
        pauseOutput()
        tap?.flush()
        written = 0
        segments.removeAll()
    }

    private func resetOutput(_ s: TrackSession) {
        pauseOutput()
        tap?.flush()
        written = 0
        segments = [Segment(writtenStart: 0, token: s.token, frameStart: s.framePos)]
        drainUntil = -1
        waitingSince = 0
        meters.clear()
        errorReported = -1
    }

    private func loop() {
        while !isReleased {
            while true {
                cmdLock.lock()
                let cmd = commands.isEmpty ? nil : commands.removeFirst()
                cmdLock.unlock()
                guard let cmd else { break }
                cmd()
            }
            do {
                try step()
            } catch {
                let tok = session?.token ?? -1
                dropAll()
                wantPlay = false
                publish()
                let msg = error.localizedDescription
                DispatchQueue.main.async { [weak self] in self?.listener?.onError(tok, msg) }
            }
        }
        dropAll()
        av.stop()
    }

    private func idle(_ ms: Int) {
        cmdLock.lock()
        if commands.isEmpty && !released { _ = cmdLock.wait(until: Date().addingTimeInterval(Double(ms) / 1000)) }
        cmdLock.unlock()
    }

    private func step() throws {
        guard let s = session, wantPlay else {
            setFlags(false, false)
            publish()
            idle(100)
            return
        }
        if let e = s.error() {
            if errorReported != s.token {
                errorReported = s.token
                wantPlay = false
                pauseOutput()
                setFlags(false, false)
                let tok = s.token, msg = e.localizedDescription
                DispatchQueue.main.async { [weak self] in self?.listener?.onError(tok, msg.isEmpty ? "Playback failed" : msg) }
            }
            idle(50)
            return
        }
        if drainUntil >= 0 {
            drain(s)
            return
        }
        let need = written == 0 || buffering ? Self.prefill : Self.chunk
        if !s.ready(need) && !s.finished() {
            let now = monoNanos()
            if waitingSince == 0 { waitingSince = now }
            if !buffering && (written == 0 || now - waitingSince > 60_000_000) {
                tap?.running = false
                setFlags(false, true)
            }
            publish()
            idle(5)
            return
        }
        waitingSince = 0
        try ensureOutput(s.sampleRate)
        try startOutput()
        setFlags(true, false)

        guard let tap else { return }
        if tap.ring.freeFrames < Self.chunk {
            publish()
            idle(5)
            return
        }
        let n = mixChunk(s)
        if n == 0 {
            if s.finished() { endOfTrack(s) } else { idle(2) }
            return
        }
        eq.configure(eqParams, sampleRate: s.sampleRate)
        eq.process(out, n)
        let m = masterGain
        rampGain(out, n, masterCur, m)
        masterCur = m
        limiter.process(out, n)
        _ = tap.ring.write(out, n * 2, abort: { false })
        written += Int64(n)
        publish()
    }

    private func endOfTrack(_ s: TrackSession) {
        if let nx = next, nx.sampleRate == s.sampleRate, nx.error() == nil {
            // Gapless: keep writing into the same output.
            segments.append(Segment(writtenStart: written, token: nx.token, frameStart: nx.framePos))
            session = nx
            next = nil
            s.release()
            return
        }
        // Let the output play out, then hand over.
        drainUntil = written
        drainToken = s.token
        drainDeadline = monoNanos() + written * 1_000_000_000 / Int64(max(1, trackRate)) + 1_000_000_000
    }

    private func drain(_ s: TrackSession) {
        let (presented, _) = self.presented()
        publish()
        if presented < drainUntil - 64 && monoNanos() < drainDeadline {
            idle(10)
            return
        }
        drainUntil = -1
        if let nx = next, nx.error() == nil {
            session = nx
            next = nil
            s.release()
            resetOutput(nx)
            let tok = nx.token
            DispatchQueue.main.async { [weak self] in self?.listener?.onTransition(tok) }
        } else {
            let tok = drainToken
            dropAll()
            wantPlay = false
            setFlags(false, false)
            publish()
            DispatchQueue.main.async { [weak self] in self?.listener?.onEnded(tok) }
        }
    }

    private func mixChunk(_ s: TrackSession) -> Int {
        if s.finished() { return 0 }
        var n = Self.chunk
        for d in s.decoders {
            if d.error != nil || d.ring.eos && d.ring.availableFrames == 0 { continue }
            n = min(n, d.ring.availableFrames)
        }
        if n <= 0 { return 0 }
        let (gains, version) = shared.with { (_mix, _mixVersion) }
        if s.mixVersion != version {
            for i in s.names.indices { s.target[i] = gains[s.names[i]] ?? 1 }
            s.mixVersion = version
        }
        out.update(repeating: 0, count: n * 2)
        let stride = MeterHistory.maxStems
        let block = MeterHistory.block
        let blocks = (n + block - 1) / block
        if meterSums.count < blocks * stride { meterSums = [Float](repeating: 0, count: blocks * stride) }
        for k in meterSums.indices { meterSums[k] = 0 }
        for (i, d) in s.decoders.enumerated() {
            let got = d.ring.read(tmp, n)
            if got < n { (tmp + got * 2).update(repeating: 0, count: (n - got) * 2) }
            let g0 = s.current[i]
            let g1 = s.target[i]
            s.current[i] = g1
            if g0 == 0 && g1 == 0 { continue }
            let step = (g1 - g0) / Float(n)
            var g = g0
            for b in 0..<blocks {
                let end = min(n, (b + 1) * block)
                var sum: Float = 0
                for f in (b * block)..<end {
                    let l = tmp[f * 2] * g
                    let r = tmp[f * 2 + 1] * g
                    out[f * 2] += l
                    out[f * 2 + 1] += r
                    sum += l * l + r * r
                    g += step
                }
                if i < stride { meterSums[b * stride + i] = (sum / Float(2 * (end - b * block))).squareRoot() }
            }
        }
        meterSums.withUnsafeBufferPointer { p in
            for b in 0..<blocks { meters.add(s.token, s.framePos + Int64(b * block), p.baseAddress! + b * stride, s.names.count) }
        }
        s.framePos += Int64(n)
        return n
    }

    private func rampGain(_ buf: UnsafeMutablePointer<Float>, _ n: Int, _ from: Float, _ to: Float) {
        if from == 1 && to == 1 { return }
        let step = (to - from) / Float(n)
        var g = from
        for f in 0..<n {
            buf[f * 2] *= g
            buf[f * 2 + 1] *= g
            g += step
        }
    }

    private static func outputLatency() -> Double {
        #if os(iOS)
        let s = AVAudioSession.sharedInstance()
        return s.outputLatency + s.ioBufferDuration
        #else
        return 0.01
        #endif
    }

    private func ensureOutput(_ rate: Int) throws {
        if node != nil && trackRate == rate { return }
        av.stop()
        if let node { av.detach(node) }
        #if os(iOS)
        try? AVAudioSession.sharedInstance().setPreferredSampleRate(Double(rate))
        #endif
        let t = OutputTap(rate: rate, frames: max(4096, rate * Self.bufferMs / 1000))
        t.setLatency(Self.outputLatency())
        guard let fmt = AVAudioFormat(standardFormatWithSampleRate: Double(rate), channels: 2) else {
            throw DecodeError("Unsupported sample rate \(rate)")
        }
        let n = AVAudioSourceNode(format: fmt) { silence, ts, frames, abl in t.render(silence, ts, frames, abl) }
        av.attach(n)
        av.connect(n, to: av.mainMixerNode, format: fmt)
        av.prepare()
        node = n
        tap = t
        trackRate = rate
        limiter.setSampleRate(rate)
    }

    private func startOutput() throws {
        guard let tap else { return }
        if !av.isRunning {
            #if os(iOS)
            try AVAudioSession.sharedInstance().setActive(true)
            #endif
            try av.start()
            tap.setLatency(Self.outputLatency())
        }
        tap.running = true
    }

    private func pauseOutput() {
        tap?.running = false
        if av.isRunning { av.pause() }
    }

    private func restartOutput() {
        tap?.setLatency(Self.outputLatency())
        // A paused engine also gets stopped by a route change; step() starts it again when it plays.
        if wantPlay, tap?.running == true, !av.isRunning { try? av.start() }
    }

    /// (frames presented since the last flush, at monoNanos).
    private func presented() -> (Int64, Int64) {
        let now = monoNanos()
        guard let tap else { return (0, now) }
        return (tap.presented(now, written), now)
    }

    private func publish() {
        let (p, nanos) = presented()
        while segments.count > 1 && segments[1].writtenStart <= p {
            segments.removeFirst()
            let tok = segments[0].token
            DispatchQueue.main.async { [weak self] in self?.listener?.onTransition(tok) }
        }
        let s = session
        let rate = s.map { $0.decoders.isEmpty ? trackRate : $0.sampleRate } ?? trackRate
        if let seg = segments.first {
            let running = playing && drainUntil < 0 || drainUntil >= 0 && p < drainUntil
            snapshot = Snapshot(token: seg.token, frame: seg.frameStart + max(0, p - seg.writtenStart), rate: rate, nanos: nanos, running: running)
        } else {
            snapshot = Snapshot(token: s?.token ?? -1, frame: s?.framePos ?? 0, rate: rate, nanos: nanos, running: false)
        }
    }

    private func setFlags(_ playing: Bool, _ buffering: Bool) {
        let changed = shared.with { () -> Bool in
            if playing == _playing && buffering == _buffering { return false }
            _playing = playing
            _buffering = buffering
            return true
        }
        if changed { DispatchQueue.main.async { [weak self] in self?.listener?.onEngineState() } }
    }
}