Cagire/src/engine/sequencer.rs

use crossbeam_channel::{bounded, Receiver, Sender, TrySendError};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use super::LinkState;
use crate::config::{MAX_BANKS, MAX_PATTERNS};
use crate::model::{ExecutionTrace, Rng, ScriptEngine, StepContext, Variables};
use crate::state::LiveKeyState;

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct PatternId {
    pub bank: usize,
    pub pattern: usize,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum PatternChange {
    Start { bank: usize, pattern: usize },
    Stop { bank: usize, pattern: usize },
}

impl PatternChange {
    pub fn pattern_id(&self) -> PatternId {
        match self {
            PatternChange::Start { bank, pattern } => PatternId {
                bank: *bank,
                pattern: *pattern,
            },
            PatternChange::Stop { bank, pattern } => PatternId {
                bank: *bank,
                pattern: *pattern,
            },
        }
    }
}

pub enum AudioCommand {
    Evaluate(String),
    Hush,
    Panic,
    LoadSamples(Vec<doux::sample::SampleEntry>),
    #[allow(dead_code)]
    ResetEngine,
}

pub enum SeqCommand {
    PatternUpdate {
        bank: usize,
        pattern: usize,
        data: PatternSnapshot,
    },
    PatternStart {
        bank: usize,
        pattern: usize,
    },
    PatternStop {
        bank: usize,
        pattern: usize,
    },
    Shutdown,
}

#[derive(Clone)]
pub struct PatternSnapshot {
    pub speed: crate::model::PatternSpeed,
    pub length: usize,
    pub steps: Vec<StepSnapshot>,
}

#[derive(Clone)]
pub struct StepSnapshot {
    pub active: bool,
    pub script: String,
    pub source: Option<usize>,
}

#[derive(Clone, Copy, Default, Debug)]
pub struct ActivePatternState {
    pub bank: usize,
    pub pattern: usize,
    pub step_index: usize,
    pub iter: usize,
}

#[derive(Clone, Default)]
pub struct SharedSequencerState {
    pub active_patterns: Vec<ActivePatternState>,
    pub step_traces: HashMap<(usize, usize, usize), ExecutionTrace>,
    pub event_count: usize,
}

pub struct SequencerSnapshot {
    pub active_patterns: Vec<ActivePatternState>,
    pub step_traces: HashMap<(usize, usize, usize), ExecutionTrace>,
    pub event_count: usize,
}

impl SequencerSnapshot {
    pub fn is_playing(&self, bank: usize, pattern: usize) -> bool {
        self.active_patterns
            .iter()
            .any(|p| p.bank == bank && p.pattern == pattern)
    }

    pub fn get_step(&self, bank: usize, pattern: usize) -> Option<usize> {
        self.active_patterns
            .iter()
            .find(|p| p.bank == bank && p.pattern == pattern)
            .map(|p| p.step_index)
    }

    pub fn get_iter(&self, bank: usize, pattern: usize) -> Option<usize> {
        self.active_patterns
            .iter()
            .find(|p| p.bank == bank && p.pattern == pattern)
            .map(|p| p.iter)
    }

    pub fn get_trace(&self, bank: usize, pattern: usize, step: usize) -> Option<&ExecutionTrace> {
        self.step_traces.get(&(bank, pattern, step))
    }
}

pub struct SequencerHandle {
    pub cmd_tx: Sender<SeqCommand>,
    pub audio_tx: Sender<AudioCommand>,
    pub audio_rx: Receiver<AudioCommand>,
    shared_state: Arc<Mutex<SharedSequencerState>>,
    thread: JoinHandle<()>,
}

impl SequencerHandle {
    pub fn snapshot(&self) -> SequencerSnapshot {
        let state = self.shared_state.lock().unwrap();
        SequencerSnapshot {
            active_patterns: state.active_patterns.clone(),
            step_traces: state.step_traces.clone(),
            event_count: state.event_count,
        }
    }

    pub fn shutdown(self) {
        let _ = self.cmd_tx.send(SeqCommand::Shutdown);
        let _ = self.thread.join();
    }
}

#[derive(Clone, Copy, Default)]
struct ActivePattern {
    bank: usize,
    pattern: usize,
    step_index: usize,
    iter: usize,
}

struct AudioState {
    prev_beat: f64,
    active_patterns: HashMap<PatternId, ActivePattern>,
    pending_starts: Vec<PatternId>,
    pending_stops: Vec<PatternId>,
}

impl AudioState {
    fn new() -> Self {
        Self {
            prev_beat: -1.0,
            active_patterns: HashMap::new(),
            pending_starts: Vec::new(),
            pending_stops: Vec::new(),
        }
    }
}

pub fn spawn_sequencer(
    link: Arc<LinkState>,
    playing: Arc<std::sync::atomic::AtomicBool>,
    variables: Variables,
    rng: Rng,
    quantum: f64,
    live_keys: Arc<LiveKeyState>,
) -> SequencerHandle {
    let (cmd_tx, cmd_rx) = bounded::<SeqCommand>(64);
    let (audio_tx, audio_rx) = bounded::<AudioCommand>(256);

    let shared_state = Arc::new(Mutex::new(SharedSequencerState::default()));
    let shared_state_clone = Arc::clone(&shared_state);
    let audio_tx_clone = audio_tx.clone();

    let thread = thread::Builder::new()
        .name("sequencer".into())
        .spawn(move || {
            sequencer_loop(
                cmd_rx,
                audio_tx_clone,
                link,
                playing,
                variables,
                rng,
                quantum,
                shared_state_clone,
                live_keys,
            );
        })
        .expect("Failed to spawn sequencer thread");

    SequencerHandle {
        cmd_tx,
        audio_tx,
        audio_rx,
        shared_state,
        thread,
    }
}

struct PatternCache {
    patterns: [[Option<PatternSnapshot>; MAX_PATTERNS]; MAX_BANKS],
}

impl PatternCache {
    fn new() -> Self {
        Self {
            patterns: std::array::from_fn(|_| std::array::from_fn(|_| None)),
        }
    }

    fn get(&self, bank: usize, pattern: usize) -> Option<&PatternSnapshot> {
        self.patterns
            .get(bank)
            .and_then(|b| b.get(pattern))
            .and_then(|p| p.as_ref())
    }

    fn set(&mut self, bank: usize, pattern: usize, data: PatternSnapshot) {
        if bank < MAX_BANKS && pattern < MAX_PATTERNS {
            self.patterns[bank][pattern] = Some(data);
        }
    }
}

impl PatternSnapshot {
    fn resolve_source(&self, index: usize) -> usize {
        let mut current = index;
        for _ in 0..self.steps.len() {
            if let Some(step) = self.steps.get(current) {
                if let Some(source) = step.source {
                    current = source;
                } else {
                    return current;
                }
            } else {
                return index;
            }
        }
        index
    }

    fn resolve_script(&self, index: usize) -> Option<&str> {
        let source_idx = self.resolve_source(index);
        self.steps.get(source_idx).map(|s| s.script.as_str())
    }
}

type StepKey = (usize, usize, usize);

struct RunsCounter {
    counts: HashMap<StepKey, usize>,
}

impl RunsCounter {
    fn new() -> Self {
        Self {
            counts: HashMap::new(),
        }
    }

    fn get_and_increment(&mut self, bank: usize, pattern: usize, step: usize) -> usize {
        let key = (bank, pattern, step);
        let count = self.counts.entry(key).or_insert(0);
        let current = *count;
        *count += 1;
        current
    }
}

#[allow(clippy::too_many_arguments)]
fn sequencer_loop(
    cmd_rx: Receiver<SeqCommand>,
    audio_tx: Sender<AudioCommand>,
    link: Arc<LinkState>,
    playing: Arc<std::sync::atomic::AtomicBool>,
    variables: Variables,
    rng: Rng,
    quantum: f64,
    shared_state: Arc<Mutex<SharedSequencerState>>,
    live_keys: Arc<LiveKeyState>,
) {
    use std::sync::atomic::Ordering;

    let script_engine = ScriptEngine::new(Arc::clone(&variables), rng);
    let mut audio_state = AudioState::new();
    let mut pattern_cache = PatternCache::new();
    let mut runs_counter = RunsCounter::new();
    let mut step_traces: HashMap<(usize, usize, usize), ExecutionTrace> = HashMap::new();
    let mut event_count: usize = 0;

    loop {
        while let Ok(cmd) = cmd_rx.try_recv() {
            match cmd {
                SeqCommand::PatternUpdate {
                    bank,
                    pattern,
                    data,
                } => {
                    pattern_cache.set(bank, pattern, data);
                }
                SeqCommand::PatternStart { bank, pattern } => {
                    let id = PatternId { bank, pattern };
                    audio_state.pending_stops.retain(|p| *p != id);
                    if !audio_state.pending_starts.contains(&id) {
                        audio_state.pending_starts.push(id);
                    }
                }
                SeqCommand::PatternStop { bank, pattern } => {
                    let id = PatternId { bank, pattern };
                    audio_state.pending_starts.retain(|p| *p != id);
                    if !audio_state.pending_stops.contains(&id) {
                        audio_state.pending_stops.push(id);
                    }
                }
                SeqCommand::Shutdown => {
                    return;
                }
            }
        }

        if !playing.load(Ordering::Relaxed) {
            thread::sleep(Duration::from_micros(500));
            continue;
        }

        let state = link.capture_app_state();
        let time = link.clock_micros();
        let beat = state.beat_at_time(time, quantum);
        let tempo = state.tempo();

        let bar = (beat / quantum).floor() as i64;
        let prev_bar = (audio_state.prev_beat / quantum).floor() as i64;
        if bar != prev_bar && audio_state.prev_beat >= 0.0 {
            for id in audio_state.pending_starts.drain(..) {
                audio_state.active_patterns.insert(
                    id,
                    ActivePattern {
                        bank: id.bank,
                        pattern: id.pattern,
                        step_index: 0,
                        iter: 0,
                    },
                );
            }
            for id in audio_state.pending_stops.drain(..) {
                audio_state.active_patterns.remove(&id);
                step_traces.retain(|&(bank, pattern, _), _| {
                    bank != id.bank || pattern != id.pattern
                });
            }
        }

        let prev_beat = audio_state.prev_beat;

        for (_id, active) in audio_state.active_patterns.iter_mut() {
            let Some(pattern) = pattern_cache.get(active.bank, active.pattern) else {
                continue;
            };

            let speed_mult = pattern.speed.multiplier();
            let beat_int = (beat * 4.0 * speed_mult).floor() as i64;
            let prev_beat_int = (prev_beat * 4.0 * speed_mult).floor() as i64;

            if beat_int != prev_beat_int && prev_beat >= 0.0 {
                let step_idx = active.step_index % pattern.length;

                if let Some(step) = pattern.steps.get(step_idx) {
                    let resolved_script = pattern.resolve_script(step_idx);
                    let has_script = resolved_script
                        .map(|s| !s.trim().is_empty())
                        .unwrap_or(false);

                    if step.active && has_script {
                        let source_idx = pattern.resolve_source(step_idx);
                        let runs =
                            runs_counter.get_and_increment(active.bank, active.pattern, source_idx);
                        let ctx = StepContext {
                            step: step_idx,
                            beat,
                            pattern: active.pattern,
                            tempo,
                            phase: beat % quantum,
                            slot: 0,
                            runs,
                            iter: active.iter,
                            speed: speed_mult,
                            fill: live_keys.fill(),
                        };
                        if let Some(script) = resolved_script {
                            let mut trace = ExecutionTrace::default();
                            if let Ok(cmds) =
                                script_engine.evaluate_with_trace(script, &ctx, &mut trace)
                            {
                                step_traces.insert(
                                    (active.bank, active.pattern, source_idx),
                                    std::mem::take(&mut trace),
                                );
                                for cmd in cmds {
                                    match audio_tx.try_send(AudioCommand::Evaluate(cmd)) {
                                        Ok(()) => {
                                            event_count += 1;
                                        }
                                        Err(TrySendError::Full(_)) => {}
                                        Err(TrySendError::Disconnected(_)) => {
                                            return;
                                        }
                                    }
                                }
                                if let Some(new_tempo) = {
                                    let mut vars = variables.lock().unwrap();
                                    vars.remove("__tempo__").and_then(|v| v.as_float().ok())
                                } {
                                    link.set_tempo(new_tempo);
                                }
                            }
                        }
                    }
                }

                let next_step = active.step_index + 1;
                if next_step >= pattern.length {
                    active.iter += 1;
                }
                active.step_index = next_step % pattern.length;
            }
        }

        {
            let mut state = shared_state.lock().unwrap();
            state.active_patterns = audio_state
                .active_patterns
                .values()
                .map(|a| ActivePatternState {
                    bank: a.bank,
                    pattern: a.pattern,
                    step_index: a.step_index,
                    iter: a.iter,
                })
                .collect();
            state.step_traces = step_traces.clone();
            state.event_count = event_count;
        }

        audio_state.prev_beat = beat;

        thread::sleep(Duration::from_micros(500));
    }
}