//! ZHILA audio engine, P0.
//!
//! A real cpal stream with a fixed voice pool, a two-bar phrase, per-vein
//! levels and a bar position. Everything that runs in the callback is
//! allocation-free. The control surface is lock-free on the audio side.

use std::f32::consts::{FRAC_PI_4, TAU};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};

use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use cpal::{BufferSize, FromSample, SampleFormat, SizedSample, Stream, StreamConfig};
use rtrb::{Consumer, Producer, RingBuffer};
use thiserror::Error;
use zhila_core::{Engine as EngineControl, MaterialSnapshot};
use zhila_protocol::{
    AudioInfo, LOOP_STEPS, MAX_BPM, MAX_GAIN, MIN_BPM, MIN_GAIN, MeterFrame, Pattern,
    STEPS_PER_BAR, VEIN_COUNT, VeinId,
};

const VOICES: usize = 32;

/// (vein index, velocity, optional (core index, face index)) for one step.
type StepTrigger = (usize, u8, Option<(usize, u8)>);

enum AudioCommand {
    Play,
    Stop,
    Strike {
        vein: u8,
        velocity: f32,
    },
    StrikeFace {
        core: String,
        face: u8,
        velocity: f32,
    },
    SetTempo(f32),
    SetGain {
        vein: u8,
        gain: f32,
    },
    SetMaterial(Arc<MaterialSnapshot>),
}

struct Levels {
    master: AtomicU32,
    veins: [AtomicU32; VEIN_COUNT],
    bar: AtomicU32,
}

impl Levels {
    fn new() -> Self {
        Self {
            master: AtomicU32::new(0.0f32.to_bits()),
            veins: std::array::from_fn(|_| AtomicU32::new(0.0f32.to_bits())),
            bar: AtomicU32::new(1.0f32.to_bits()),
        }
    }

    fn store(&self, frame: &MeterFrame) {
        self.master.store(frame.master.to_bits(), Ordering::Relaxed);
        self.bar.store(frame.bar.to_bits(), Ordering::Relaxed);
        for (level, slot) in frame.veins.iter().zip(self.veins.iter()) {
            slot.store(level.to_bits(), Ordering::Relaxed);
        }
    }

    fn load(&self) -> MeterFrame {
        let master = f32::from_bits(self.master.load(Ordering::Relaxed));
        let bar = f32::from_bits(self.bar.load(Ordering::Relaxed));
        let mut veins = [0.0; VEIN_COUNT];
        for (level, slot) in veins.iter_mut().zip(self.veins.iter()) {
            *level = f32::from_bits(slot.load(Ordering::Relaxed));
        }
        MeterFrame { bar, master, veins }
    }
}

struct SampleVoice {
    core_index: usize,
    start: usize,
    end: usize,
    position: f64,
    rate: f64,
}

struct Voice {
    active: bool,
    vein: usize,
    phase: f32,
    freq: f32,
    t: f32,
    attack: f32,
    decay: f32,
    velocity: f32,
    gain_l: f32,
    gain_r: f32,
    lp: f32,
    hp: f32,
    prev: f32,
    rng: u32,
    sample: Option<SampleVoice>,
}

impl Voice {
    fn inactive() -> Self {
        Self {
            active: false,
            vein: 0,
            phase: 0.0,
            freq: 0.0,
            t: 0.0,
            attack: 0.01,
            decay: 4.0,
            velocity: 0.0,
            gain_l: 0.707,
            gain_r: 0.707,
            lp: 0.0,
            hp: 0.0,
            prev: 0.0,
            rng: 0x9e37_79b9,
            sample: None,
        }
    }

    fn trigger(&mut self, vein: VeinId, velocity: f32, seed: u32) {
        self.active = true;
        self.vein = vein.index();
        self.phase = 0.0;
        self.t = 0.0;
        self.velocity = velocity.clamp(0.0, 1.0);
        self.rng = seed | 1;
        self.lp = 0.0;
        self.hp = 0.0;
        self.prev = 0.0;
        self.sample = None;
        let (freq, attack, decay) = match vein {
            VeinId::Bass => (55.0, 0.004, 3.2),
            VeinId::Crack => (0.0, 0.001, 38.0),
            VeinId::Voice => (220.0, 0.09, 1.1),
            VeinId::Noise => (0.0, 0.004, 5.0),
            VeinId::Silt => (1318.5, 0.002, 14.0),
        };
        self.freq = freq;
        self.attack = attack;
        self.decay = decay;
        let (gain_l, gain_r) = pan_for(vein);
        self.gain_l = gain_l;
        self.gain_r = gain_r;
    }

    fn trigger_sample(
        &mut self,
        vein: VeinId,
        velocity: f32,
        core_index: usize,
        start: usize,
        end: usize,
        rate: f64,
    ) {
        self.active = true;
        self.vein = vein.index();
        self.phase = 0.0;
        self.freq = 0.0;
        self.t = 0.0;
        self.velocity = velocity.clamp(0.0, 1.0);
        self.attack = 0.004;
        self.decay = 4.0;
        self.lp = 0.0;
        self.hp = 0.0;
        self.prev = 0.0;
        let (gain_l, gain_r) = pan_for(vein);
        self.gain_l = gain_l;
        self.gain_r = gain_r;
        self.sample = Some(SampleVoice {
            core_index,
            start,
            end,
            position: start as f64,
            rate,
        });
    }

    fn noise(&mut self) -> f32 {
        self.rng = self.rng.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
        (self.rng as f32 / u32::MAX as f32) * 2.0 - 1.0
    }

    fn next(&mut self, sample_rate: f32, material: Option<&MaterialSnapshot>) -> f32 {
        self.t += 1.0 / sample_rate;
        if let Some(voice) = self.sample.as_mut() {
            let Some(core) = material.and_then(|material| material.cores.get(voice.core_index))
            else {
                self.active = false;
                return 0.0;
            };
            if voice.position >= voice.end as f64 {
                self.active = false;
                return 0.0;
            }
            let index = voice.position.floor() as usize;
            let fraction = (voice.position - index as f64) as f32;
            let first = core.samples.get(index).copied().unwrap_or(0.0);
            let second = core.samples.get(index + 1).copied().unwrap_or(first);
            let raw = first + (second - first) * fraction;
            let remaining =
                ((voice.end as f64 - voice.position) / voice.rate / sample_rate as f64) as f32;
            let fade_in = (self.t / 0.004).min(1.0);
            let fade_out = (remaining / 0.004).min(1.0);
            voice.position += voice.rate;
            if voice.position >= voice.end as f64 {
                self.active = false;
            }
            let _ = voice.start;
            return raw * fade_in.min(fade_out) * self.velocity;
        }

        let attack = (self.t / self.attack).min(1.0);
        let env = attack * (-self.t * self.decay).exp();
        if env < 0.0002 || self.t > 8.0 {
            self.active = false;
            return 0.0;
        }
        let sample = match self.vein {
            0 => {
                let fundamental = (TAU * self.phase).sin();
                let sub = (TAU * self.phase * 0.5).sin();
                fundamental + sub * 0.45
            }
            1 => {
                let white = self.noise();
                let hp = white - self.prev + 0.94 * self.hp;
                self.prev = white;
                self.hp = hp;
                hp * 0.7
            }
            2 => {
                let fundamental = (TAU * self.phase).sin();
                let overtone = (TAU * self.phase * 2.0).sin();
                fundamental + overtone * 0.18
            }
            3 => {
                let white = self.noise();
                self.lp += (white - self.lp) * 0.06;
                self.lp
            }
            _ => (TAU * self.phase).sin(),
        };
        if self.freq > 0.0 {
            self.phase += self.freq / sample_rate;
            if self.phase >= 1.0 {
                self.phase -= 1.0;
            }
        }
        sample * env * self.velocity
    }
}

fn pan_for(vein: VeinId) -> (f32, f32) {
    let pan = match vein {
        VeinId::Bass => 0.0,
        VeinId::Crack => 0.35,
        VeinId::Voice => -0.25,
        VeinId::Noise => 0.55,
        VeinId::Silt => -0.45,
    };
    let angle = (pan + 1.0) * FRAC_PI_4;
    (angle.cos(), angle.sin())
}

struct AudioState {
    rx: Consumer<AudioCommand>,
    sample_rate: f32,
    channels: usize,
    bpm: f32,
    playing: bool,
    step: usize,
    pending: f64,
    fallback: Pattern,
    material: Option<Arc<MaterialSnapshot>>,
    voices: [Voice; VOICES],
    gains: [f32; VEIN_COUNT],
    master_env: f32,
    vein_env: [f32; VEIN_COUNT],
    bar: f32,
    rng: u32,
}

impl AudioState {
    fn new(sample_rate: u32, channels: u16, pattern: &Pattern, rx: Consumer<AudioCommand>) -> Self {
        Self {
            rx,
            sample_rate: sample_rate as f32,
            channels: channels as usize,
            bpm: 128.0,
            playing: false,
            step: 0,
            pending: 0.0,
            fallback: pattern.clone(),
            material: None,
            voices: std::array::from_fn(|_| Voice::inactive()),
            gains: zhila_core::INITIAL_GAINS,
            master_env: 0.0,
            vein_env: [0.0; VEIN_COUNT],
            bar: 1.0,
            rng: 0x51ed_2701,
        }
    }

    fn step_samples(&self) -> f64 {
        self.sample_rate as f64 * 60.0 / self.bpm as f64 / 2.0
    }

    fn handle(&mut self, command: AudioCommand) {
        match command {
            AudioCommand::Play => {
                self.playing = true;
                self.step = 0;
                self.pending = 0.0;
                self.bar = 1.0;
            }
            AudioCommand::Stop => {
                self.playing = false;
                self.step = 0;
                self.pending = 0.0;
                self.bar = 1.0;
                self.master_env = 0.0;
                self.vein_env = [0.0; VEIN_COUNT];
                for voice in &mut self.voices {
                    voice.active = false;
                }
            }
            AudioCommand::Strike { vein, velocity } => {
                let Some(vein) = VeinId::from_index(vein as usize) else {
                    return;
                };
                self.trigger(vein, velocity);
            }
            AudioCommand::StrikeFace {
                core,
                face,
                velocity,
            } => {
                let Some(material) = self.material.as_ref().cloned() else {
                    return;
                };
                let Some(core_index) = material
                    .cores
                    .iter()
                    .position(|candidate| candidate.id == core)
                else {
                    return;
                };
                let vein_index = material
                    .pattern
                    .rows
                    .iter()
                    .find(|row| {
                        row.face.as_ref().is_some_and(|reference| {
                            reference.core == core && reference.face == face
                        })
                    })
                    .map(|row| row.vein.index())
                    .unwrap_or(0);
                self.trigger_sample(core_index, face, vein_index, velocity);
            }
            AudioCommand::SetTempo(bpm) => {
                self.bpm = bpm.clamp(MIN_BPM, MAX_BPM);
            }
            AudioCommand::SetGain { vein, gain } => {
                if let Some(index) = VeinId::from_index(vein as usize) {
                    self.gains[index.index()] = gain.clamp(MIN_GAIN, MAX_GAIN);
                }
            }
            AudioCommand::SetMaterial(material) => {
                self.material = Some(material);
            }
        }
    }

    fn free_voice(&self) -> usize {
        self.voices
            .iter()
            .position(|voice| !voice.active)
            .unwrap_or_else(|| {
                self.voices
                    .iter()
                    .enumerate()
                    .max_by(|a, b| a.1.t.total_cmp(&b.1.t))
                    .map(|(index, _)| index)
                    .unwrap_or(0)
            })
    }

    fn trigger(&mut self, vein: VeinId, velocity: f32) {
        if velocity <= 0.0 {
            return;
        }
        let index = self.free_voice();
        self.rng = self.rng.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
        self.voices[index].trigger(vein, velocity, self.rng);
    }

    fn trigger_sample(
        &mut self,
        core_index: usize,
        face_index: u8,
        vein_index: usize,
        velocity: f32,
    ) {
        if velocity <= 0.0 {
            return;
        }
        let Some(material) = self.material.as_ref().cloned() else {
            return;
        };
        let Some(core) = material.cores.get(core_index) else {
            return;
        };
        let Some(face) = core.faces.iter().find(|face| face.index == face_index) else {
            return;
        };
        let start = face.start as usize;
        let end = face.end as usize;
        if start >= end || end > core.samples.len() {
            return;
        }
        let Some(vein) = VeinId::from_index(vein_index) else {
            return;
        };
        let rate = core.sample_rate as f64 / self.sample_rate as f64;
        let index = self.free_voice();
        self.voices[index].trigger_sample(vein, velocity, core_index, start, end, rate);
    }

    fn trigger_step(&mut self) {
        let step = self.step;
        let (plan, count) = {
            let material = self.material.as_ref().cloned();
            let pattern = material
                .as_ref()
                .map(|material| &material.pattern)
                .unwrap_or(&self.fallback);
            let mut plan: [StepTrigger; 8] = [(0, 0, None); 8];
            let mut count = 0;
            for row in &pattern.rows {
                if count >= plan.len() {
                    break;
                }
                let velocity = row.velocities.get(step).copied().unwrap_or(0);
                if velocity == 0 {
                    continue;
                }
                let face = row.face.as_ref().and_then(|reference| {
                    material.as_ref().and_then(|material| {
                        material
                            .cores
                            .iter()
                            .position(|core| core.id == reference.core)
                            .map(|core_index| (core_index, reference.face))
                    })
                });
                plan[count] = (row.vein.index(), velocity, face);
                count += 1;
            }
            (plan, count)
        };
        for (vein_index, velocity, face) in plan.into_iter().take(count) {
            if let Some((core_index, face_index)) = face {
                self.trigger_sample(core_index, face_index, vein_index, velocity as f32 / 127.0);
            } else if let Some(vein) = VeinId::from_index(vein_index) {
                self.trigger(vein, velocity as f32 / 127.0);
            }
        }
        self.bar = step as f32 / STEPS_PER_BAR as f32 + 1.0;
        self.step = (step + 1) % LOOP_STEPS;
    }

    fn render<T>(&mut self, data: &mut [T], levels: &Levels)
    where
        T: SizedSample + FromSample<f32>,
    {
        while let Ok(command) = self.rx.pop() {
            self.handle(command);
        }

        let material = self.material.as_ref().cloned();
        for frame in data.chunks_mut(self.channels) {
            if self.playing {
                if self.pending <= 0.0 {
                    self.trigger_step();
                    self.pending += self.step_samples();
                }
                self.pending -= 1.0;
            }

            let mut left = 0.0f32;
            let mut right = 0.0f32;
            let mut sums = [0.0f32; VEIN_COUNT];
            for voice in &mut self.voices {
                if !voice.active {
                    continue;
                }
                let sample = voice.next(self.sample_rate, material.as_deref());
                let index = voice.vein;
                let amplitude = sample * self.gains[index];
                left += amplitude * voice.gain_l;
                right += amplitude * voice.gain_r;
                sums[index] += amplitude.abs();
            }

            let peak = (left.abs() + right.abs()) * 0.5;
            self.master_env += (peak - self.master_env) * 0.0015;
            for (env, sum) in self.vein_env.iter_mut().zip(sums.iter()) {
                *env += (*sum - *env) * 0.002;
            }

            let out_l = (left * 0.9).tanh() * 0.8;
            let out_r = (right * 0.9).tanh() * 0.8;
            for (channel, sample) in frame.iter_mut().enumerate() {
                let value = match (self.channels, channel) {
                    (1, _) => (out_l + out_r) * 0.5,
                    (_, 0) => out_l,
                    (_, 1) => out_r,
                    _ => (out_l + out_r) * 0.5,
                };
                *sample = T::from_sample(value);
            }
        }

        levels.store(&MeterFrame {
            bar: self.bar,
            master: self.master_env,
            veins: self.vein_env,
        });
    }
}

/// The send side of the engine. Cheap, lock-free for the audio thread and
/// safe to share with the session and the WebSocket daemon.
pub struct Control {
    tx: Mutex<Producer<AudioCommand>>,
    levels: Arc<Levels>,
    info: AudioInfo,
}

impl EngineControl for Control {
    fn play(&self) {
        self.send(AudioCommand::Play);
    }

    fn stop(&self) {
        self.send(AudioCommand::Stop);
    }

    fn strike(&self, vein: VeinId, velocity: f32) {
        self.send(AudioCommand::Strike {
            vein: vein.index() as u8,
            velocity,
        });
    }

    fn strike_face(&self, core: &str, face: u8, velocity: f32) {
        self.send(AudioCommand::StrikeFace {
            core: core.to_string(),
            face,
            velocity,
        });
    }

    fn set_tempo(&self, bpm: f32) {
        self.send(AudioCommand::SetTempo(bpm));
    }

    fn set_gain(&self, vein: VeinId, gain: f32) {
        self.send(AudioCommand::SetGain {
            vein: vein.index() as u8,
            gain,
        });
    }

    fn set_material(&self, material: Arc<MaterialSnapshot>) {
        self.send(AudioCommand::SetMaterial(material));
    }

    fn info(&self) -> AudioInfo {
        self.info.clone()
    }

    fn meter(&self) -> MeterFrame {
        self.levels.load()
    }
}

impl Control {
    fn send(&self, command: AudioCommand) {
        if let Ok(mut tx) = self.tx.lock() {
            let _ = tx.push(command);
        }
    }
}

#[derive(Debug, Error)]
pub enum EngineError {
    #[error("no default audio output device")]
    NoDevice,
    #[error("audio backend: {0}")]
    Backend(#[from] cpal::Error),
    #[error("unsupported sample format: {0:?}")]
    UnsupportedFormat(SampleFormat),
}

struct Started {
    stream: Stream,
    control: Arc<Control>,
}

/// Keeps the cpal stream alive. The control handle can be shared freely; the
/// stream itself must stay on the thread that created it.
pub struct AudioEngine {
    _stream: Stream,
    control: Arc<Control>,
}

impl AudioEngine {
    pub fn start(pattern: &Pattern) -> Result<Self, EngineError> {
        let started = match Self::try_start(pattern, BufferSize::Fixed(256)) {
            Ok(started) => started,
            Err(EngineError::Backend(_)) => Self::try_start(pattern, BufferSize::Default)?,
            Err(error) => return Err(error),
        };
        Ok(Self {
            _stream: started.stream,
            control: started.control,
        })
    }

    pub fn control(&self) -> Arc<Control> {
        self.control.clone()
    }

    pub fn info(&self) -> AudioInfo {
        self.control.info.clone()
    }

    fn try_start(pattern: &Pattern, buffer: BufferSize) -> Result<Started, EngineError> {
        let host = cpal::default_host();
        let device = host.default_output_device().ok_or(EngineError::NoDevice)?;
        let supported = device.default_output_config()?;
        let format = supported.sample_format();
        let mut config: StreamConfig = supported.into();
        config.buffer_size = buffer;
        let sample_rate = config.sample_rate;
        let channels = config.channels;
        let device_name = device.to_string();
        let (producer, consumer) = RingBuffer::new(256);
        let levels = Arc::new(Levels::new());
        let state = AudioState::new(sample_rate, channels, pattern, consumer);
        let stream = match format {
            SampleFormat::F32 => build_stream::<f32>(&device, config, state, levels.clone())?,
            SampleFormat::I16 => build_stream::<i16>(&device, config, state, levels.clone())?,
            other => return Err(EngineError::UnsupportedFormat(other)),
        };
        let info = AudioInfo {
            running: true,
            device: device_name,
            sample_rate,
            buffer_frames: match buffer {
                BufferSize::Fixed(frames) => frames,
                BufferSize::Default => 0,
            },
            channels,
        };
        let control = Arc::new(Control {
            tx: Mutex::new(producer),
            levels,
            info,
        });
        Ok(Started { stream, control })
    }
}

fn build_stream<T>(
    device: &cpal::Device,
    config: StreamConfig,
    mut state: AudioState,
    levels: Arc<Levels>,
) -> Result<Stream, EngineError>
where
    T: SizedSample + FromSample<f32> + 'static,
{
    let stream = device.build_output_stream(
        config,
        move |data: &mut [T], _: &cpal::OutputCallbackInfo| {
            state.render(data, &levels);
        },
        |error| eprintln!("zhila audio: {error}"),
        None,
    )?;
    stream.play()?;
    Ok(stream)
}

/// A silent engine for headless protocol tests and machines without an output
/// device. It reports itself as offline so the surface never pretends to play.
pub struct NullEngine {
    info: AudioInfo,
}

impl Default for NullEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl NullEngine {
    pub fn new() -> Self {
        Self {
            info: AudioInfo::offline(),
        }
    }
}

impl EngineControl for NullEngine {
    fn play(&self) {}
    fn stop(&self) {}
    fn strike(&self, _vein: VeinId, _velocity: f32) {}
    fn strike_face(&self, _core: &str, _face: u8, _velocity: f32) {}
    fn set_tempo(&self, _bpm: f32) {}
    fn set_gain(&self, _vein: VeinId, _gain: f32) {}
    fn set_material(&self, _material: Arc<MaterialSnapshot>) {}
    fn info(&self) -> AudioInfo {
        self.info.clone()
    }
    fn meter(&self) -> MeterFrame {
        MeterFrame::silent()
    }
}