import Foundation

/// ISO octave centres for the 10-band graphic EQ.
let EQ_BANDS: [Float] = [31.25, 62.5, 125, 250, 500, 1000, 2000, 4000, 8000, 16000]
let EQ_LABELS = ["31", "62", "125", "250", "500", "1K", "2K", "4K", "8K", "16K"]
let EQ_MAX_DB: Float = 12
let EQ_Q: Float = 1.41 // one octave

/// Equalizer APO's filter types; [apo] is the name its config files use.
enum FilterType: String, Codable, CaseIterable {
    case Peak, LowShelf, HighShelf, LowPass, HighPass, BandPass, Notch, AllPass

    var apo: String {
        switch self {
        case .Peak: "PK"
        case .LowShelf: "LSC"
        case .HighShelf: "HSC"
        case .LowPass: "LPQ"
        case .HighPass: "HPQ"
        case .BandPass: "BP"
        case .Notch: "NO"
        case .AllPass: "AP"
        }
    }

    var label: String {
        switch self {
        case .Peak: "Peak"
        case .LowShelf: "Low Shelf"
        case .HighShelf: "High Shelf"
        case .LowPass: "Low Pass"
        case .HighPass: "High Pass"
        case .BandPass: "Band Pass"
        case .Notch: "Notch"
        case .AllPass: "All Pass"
        }
    }

    var hasGain: Bool { self == .Peak || self == .LowShelf || self == .HighShelf }
    var hasSlope: Bool { self == .LowPass || self == .HighPass }
}

enum EqChannel: String, Codable, CaseIterable {
    case Both, Left, Right
    var label: String { rawValue }
}

/// Slopes offered for low and high pass, in dB per octave. 12 uses the filter's Q; steeper ones are Butterworth.
let EQ_SLOPES = [6, 12, 24, 36, 48]

/// One parametric filter, as in Equalizer APO. Gain only matters for peaks and shelves.
struct EqFilter: Codable, Hashable, Identifiable {
    var id: String
    var type: FilterType = .Peak
    var freq: Float = 1000
    var gain: Float = 0
    var q: Float = 0.707
    var slope = 12
    var channel: EqChannel = .Both
    var on = true

    init(id: String, type: FilterType = .Peak, freq: Float = 1000, gain: Float = 0, q: Float = 0.707,
         slope: Int = 12, channel: EqChannel = .Both, on: Bool = true) {
        self.id = id; self.type = type; self.freq = freq; self.gain = gain; self.q = q
        self.slope = slope; self.channel = channel; self.on = on
    }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        id = c.str("id") ?? UUID().uuidString
        type = c.str("type").flatMap(FilterType.init(rawValue:)) ?? .Peak
        freq = Float(c.double("freq") ?? 1000)
        gain = Float(c.double("gain") ?? 0)
        q = Float(c.double("q") ?? 0.707)
        slope = c.int("slope") ?? 12
        channel = c.str("channel").flatMap(EqChannel.init(rawValue:)) ?? .Both
        on = c.bool("on") ?? true
    }

    /// Whether it changes the sound at all (a 0 dB peak or shelf doesn't).
    var active: Bool { on && (!type.hasGain || abs(gain) >= 0.01) }

    /// Its biquad sections at [fs]: [b0, b1, b2, a1, a2] each, normalised by a0.
    func sections(_ fs: Double) -> [[Double]] {
        guard active else { return [] }
        let f0 = min(max(Double(freq), 1), fs * 0.49)
        let q = min(max(Double(self.q), 0.025), 100)
        let g = Double(gain)
        switch type {
        case .LowPass, .HighPass:
            let lp = type == .LowPass
            if slope <= 6 { return [firstOrder(fs, f0, lp)] }
            if slope <= 12 { return [Biquad.design(type, fs, f0, 0, q)] }
            // Butterworth of order slope/6, as cascaded sections (plus a first-order one when odd).
            let n = slope / 6
            var out = (1...(n / 2)).map { k in Biquad.design(type, fs, f0, 0, 1 / (2 * cos(Double(2 * k - 1) * .pi / Double(2 * n)))) }
            if n % 2 == 1 { out.append(firstOrder(fs, f0, lp)) }
            return out
        default:
            return [Biquad.design(type, fs, f0, g, q)]
        }
    }

    private func firstOrder(_ fs: Double, _ f0: Double, _ lowPass: Bool) -> [Double] {
        let k = tan(.pi * f0 / fs)
        let a1 = (k - 1) / (k + 1)
        return lowPass ? [k / (k + 1), k / (k + 1), 0, a1, 0] : [1 / (k + 1), -1 / (k + 1), 0, a1, 0]
    }

    func applies(left: Bool) -> Bool { channel == .Both || (channel == .Left) == left }

    /// The 10-band graphic EQ as peaking filters, ahead of the parametric ones.
    static func graphic(_ gains: [Float]) -> [EqFilter] {
        gains.enumerated().compactMap { i, g in
            abs(g) < 0.01 || i >= EQ_BANDS.count ? nil : EqFilter(id: "band\(i)", type: .Peak, freq: EQ_BANDS[i], gain: g, q: EQ_Q)
        }
    }
}

/// What the engine needs from an EQ profile; nil means bypass.
struct EqParams: Equatable {
    var filters: [EqFilter]
    var preampDb: Float
}

/// A biquad section, transposed direct form II, double state per channel.
final class Biquad {
    private var b0 = 1.0, b1 = 0.0, b2 = 0.0, a1 = 0.0, a2 = 0.0
    private var zl1 = 0.0, zl2 = 0.0, zr1 = 0.0, zr2 = 0.0

    func set(_ c: [Double]) {
        b0 = c[0]; b1 = c[1]; b2 = c[2]; a1 = c[3]; a2 = c[4]
    }

    func reset() { zl1 = 0; zl2 = 0; zr1 = 0; zr2 = 0 }

    /// [buf] is interleaved stereo.
    func process(_ buf: UnsafeMutablePointer<Float>, _ frames: Int, left: Bool = true, right: Bool = true) {
        var l1 = zl1, l2 = zl2, r1 = zr1, r2 = zr2
        for i in 0..<frames {
            if left {
                let xl = Double(buf[i * 2])
                let yl = b0 * xl + l1
                l1 = b1 * xl - a1 * yl + l2
                l2 = b2 * xl - a2 * yl
                buf[i * 2] = Float(yl)
            }
            if right {
                let xr = Double(buf[i * 2 + 1])
                let yr = b0 * xr + r1
                r1 = b1 * xr - a1 * yr + r2
                r2 = b2 * xr - a2 * yr
                buf[i * 2 + 1] = Float(yr)
            }
        }
        // Flush denormals so a silent tail doesn't cost CPU.
        zl1 = abs(l1) < 1e-25 ? 0 : l1; zl2 = abs(l2) < 1e-25 ? 0 : l2
        zr1 = abs(r1) < 1e-25 ? 0 : r1; zr2 = abs(r2) < 1e-25 ? 0 : r2
    }

    /// RBJ Audio EQ Cookbook coefficients, [b0, b1, b2, a1, a2] normalised by a0.
    static func design(_ type: FilterType, _ fs: Double, _ f0: Double, _ gainDb: Double, _ q: Double) -> [Double] {
        let a = pow(10, gainDb / 40)
        let w0 = 2 * .pi * min(f0, fs * 0.49) / fs
        let alpha = sin(w0) / (2 * q)
        let cw = cos(w0)
        let sa = 2 * a.squareRoot() * alpha
        let c: [Double]
        switch type {
        case .Peak: c = [1 + alpha * a, -2 * cw, 1 - alpha * a, 1 + alpha / a, -2 * cw, 1 - alpha / a]
        case .LowShelf: c = [
            a * ((a + 1) - (a - 1) * cw + sa), 2 * a * ((a - 1) - (a + 1) * cw), a * ((a + 1) - (a - 1) * cw - sa),
            (a + 1) + (a - 1) * cw + sa, -2 * ((a - 1) + (a + 1) * cw), (a + 1) + (a - 1) * cw - sa,
        ]
        case .HighShelf: c = [
            a * ((a + 1) + (a - 1) * cw + sa), -2 * a * ((a - 1) + (a + 1) * cw), a * ((a + 1) + (a - 1) * cw - sa),
            (a + 1) - (a - 1) * cw + sa, 2 * ((a - 1) - (a + 1) * cw), (a + 1) - (a - 1) * cw - sa,
        ]
        case .LowPass: c = [(1 - cw) / 2, 1 - cw, (1 - cw) / 2, 1 + alpha, -2 * cw, 1 - alpha]
        case .HighPass: c = [(1 + cw) / 2, -(1 + cw), (1 + cw) / 2, 1 + alpha, -2 * cw, 1 - alpha]
        case .BandPass: c = [alpha, 0, -alpha, 1 + alpha, -2 * cw, 1 - alpha]
        case .Notch: c = [1, -2 * cw, 1, 1 + alpha, -2 * cw, 1 - alpha]
        case .AllPass: c = [1 - alpha, -2 * cw, 1 + alpha, 1 + alpha, -2 * cw, 1 - alpha]
        }
        let a0 = c[3]
        return [c[0] / a0, c[1] / a0, c[2] / a0, c[4] / a0, c[5] / a0]
    }

    /// |H| in dB at frequency f.
    static func magnitudeDb(_ c: [Double], _ fs: Double, _ f: Double) -> Double {
        let w = 2 * .pi * f / fs
        let cw = cos(w), sw = sin(w), c2 = cos(2 * w), s2 = sin(2 * w)
        let nr = c[0] + c[1] * cw + c[2] * c2
        let ni = -(c[1] * sw + c[2] * s2)
        let dr = 1 + c[3] * cw + c[4] * c2
        let di = -(c[3] * sw + c[4] * s2)
        return 10 * log10(max(1e-30, nr * nr + ni * ni) / max(1e-30, dr * dr + di * di))
    }
}

enum EqMath {
    /// Combined response of [filters] on one channel, in dB at each of [freqs].
    static func response(_ filters: [EqFilter], _ freqs: [Double], left: Bool = true, fs: Double = 48000) -> [Double] {
        let coeffs = filters.filter { $0.applies(left: left) }.flatMap { $0.sections(fs) }
        return freqs.map { f in coeffs.reduce(0) { $0 + Biquad.magnitudeDb($1, fs, f) } }
    }

    static func logFreqs(_ n: Int, lo: Double = 20, hi: Double = 20000) -> [Double] {
        (0..<n).map { lo * pow(hi / lo, Double($0) / (Double(n) - 1)) }
    }

    /// Anti-clipping preamp: the negative of the curve's highest point (manual preamp included,
    /// louder channel), so a full-scale signal at the loudest frequency still peaks at 0 dBFS.
    static func autoPreampDb(_ filters: [EqFilter], manualDb: Float = 0) -> Float {
        let f = logFreqs(512)
        let peak = max(response(filters, f, left: true).max() ?? 0, response(filters, f, left: false).max() ?? 0)
        return -Float(max(0, peak + Double(manualDb)))
    }
}

/// Runs on the render thread; [configure] is called there too, between buffers.
final class Equalizer {
    private struct Stage { let filter: Biquad; let left: Bool; let right: Bool }
    private struct Shape: Equatable { let id: String; let channel: EqChannel; let type: FilterType }

    private var stages: [Stage] = []
    private var shape: [Shape] = []
    private var fs = 0
    private var params: EqParams?
    private var preamp: Float = 1
    private(set) var enabled = false

    func configure(_ p: EqParams?, sampleRate: Int) {
        enabled = p != nil
        guard let p, p != params || sampleRate != fs else { return }
        let rateChanged = sampleRate != fs
        fs = sampleRate
        params = p
        preamp = dbToGain(p.preampDb)
        let designed = p.filters.flatMap { f in f.sections(Double(fs)).map { ($0, f) } }
        // Coefficient changes (a slider drag) keep the filters' state; a different layout starts fresh.
        let newShape = designed.map { Shape(id: $0.1.id, channel: $0.1.channel, type: $0.1.type) }
        if rateChanged || newShape != shape {
            stages = designed.map { Stage(filter: Biquad(), left: $0.1.channel != .Right, right: $0.1.channel != .Left) }
            shape = newShape
        }
        for (i, d) in designed.enumerated() { stages[i].filter.set(d.0) }
    }

    func process(_ buf: UnsafeMutablePointer<Float>, _ frames: Int) {
        guard enabled else { return }
        if preamp != 1 { for i in 0..<(frames * 2) { buf[i] *= preamp } }
        for s in stages { s.filter.process(buf, frames, left: s.left, right: s.right) }
    }
}

/// Peak limiter on the master bus: stems pushed past unity or an EQ boost never clip.
/// Instant attack, ~120 ms release, ceiling just under full scale.
final class Limiter {
    private let ceiling: Float
    private var gain: Float = 1
    private var release: Float = 0

    init(ceiling: Float = 0.985) { self.ceiling = ceiling }

    func setSampleRate(_ fs: Int) { release = 1 - exp(-1 / (0.12 * Float(fs))) }

    func process(_ buf: UnsafeMutablePointer<Float>, _ frames: Int) {
        var g = gain
        for i in 0..<frames {
            let l = buf[i * 2], r = buf[i * 2 + 1]
            let peak = max(abs(l), abs(r))
            let need: Float = peak * g > ceiling ? ceiling / peak : 1
            g = need < g ? need : g + (1 - g) * release
            if peak * g > ceiling { g = ceiling / peak }
            buf[i * 2] = l * g
            buf[i * 2 + 1] = r * g
        }
        gain = g
    }
}

func dbToGain(_ db: Float) -> Float { pow(10, db / 20) }
func gainToDb(_ g: Float) -> Float { g <= 1e-6 ? -120 : 20 * log10(g) }

/// Master slider taper (same cubic curve as the desktop app).
func masterTaper(_ v: Float) -> Float { let x = min(max(v, 0), 1); return x * x * x }

/// Stem slider: 0..1 with unity at 0.8 and +3.9 dB at the top.
func stemTaper(_ v: Float) -> Float { let x = min(max(v, 0), 1) / 0.8; return x * x }
func stemUntaper(_ g: Float) -> Float { min(max(max(g, 0).squareRoot() * 0.8, 0), 1) }