Lightningbeam/daw-backend/src/audio/node_graph/nodes/voice_allocator.rs

use crate::audio::midi::MidiEvent;
use crate::audio::node_graph::{AudioNode, AudioGraph, NodeCategory, NodePort, Parameter, ParameterUnit, SignalType};

const PARAM_VOICE_COUNT: u32 = 0;
const MAX_VOICES: usize = 16; // Maximum allowed voices
const DEFAULT_VOICES: usize = 8;

/// Voice state for voice allocation
#[derive(Clone)]
struct VoiceState {
    active: bool,
    releasing: bool, // Note-off received, still processing (e.g. ADSR release)
    note: u8,
    note_channel: u8, // MIDI channel this voice was allocated on (0 = global/unset)
    age: u32, // For voice stealing
    pending_events: Vec<MidiEvent>, // MIDI events to send to this voice
}

impl VoiceState {
    fn new() -> Self {
        Self {
            active: false,
            releasing: false,
            note: 0,
            note_channel: 0,
            age: 0,
            pending_events: Vec::new(),
        }
    }
}

/// VoiceAllocatorNode - A group node that creates N polyphonic instances of its internal graph
///
/// This node acts as a container for a "voice template" graph. At runtime, it creates
/// N instances of that graph (one per voice) and routes MIDI note events to them.
/// All voice outputs are mixed together into a single output.
pub struct VoiceAllocatorNode {
    name: String,

    /// The template graph (edited by user via UI)
    template_graph: AudioGraph,

    /// Runtime voice instances (clones of template)
    voice_instances: Vec<AudioGraph>,

    /// Voice allocation state
    voices: [VoiceState; MAX_VOICES],

    /// Number of active voices (configurable parameter)
    voice_count: usize,

    /// Mix buffer for combining voice outputs
    mix_buffer: Vec<f32>,

    inputs: Vec<NodePort>,
    outputs: Vec<NodePort>,
    parameters: Vec<Parameter>,
}

impl VoiceAllocatorNode {
    pub fn new(name: impl Into<String>, sample_rate: u32, buffer_size: usize) -> Self {
        let name = name.into();

        // MIDI input for receiving note events
        let inputs = vec![
            NodePort::new("MIDI In", SignalType::Midi, 0),
        ];

        // Single mixed audio output
        let outputs = vec![
            NodePort::new("Mixed Out", SignalType::Audio, 0),
        ];

        // Voice count parameter
        let parameters = vec![
            Parameter::new(PARAM_VOICE_COUNT, "Voices", 1.0, MAX_VOICES as f32, DEFAULT_VOICES as f32, ParameterUnit::Generic),
        ];

        // Create template graph with default TemplateInput and TemplateOutput nodes
        let mut template_graph = AudioGraph::new(sample_rate, buffer_size);
        {
            use super::template_io::{TemplateInputNode, TemplateOutputNode};
            let input_node = Box::new(TemplateInputNode::new("Template Input"));
            let output_node = Box::new(TemplateOutputNode::new("Template Output"));
            let input_idx = template_graph.add_node(input_node);
            let output_idx = template_graph.add_node(output_node);
            template_graph.set_node_position(input_idx, -200.0, 0.0);
            template_graph.set_node_position(output_idx, 200.0, 0.0);
            template_graph.set_midi_target(input_idx, true);
            template_graph.set_output_node(Some(output_idx));
        }

        // Create voice instances (initially empty clones of template)
        let voice_instances: Vec<AudioGraph> = (0..MAX_VOICES)
            .map(|_| AudioGraph::new(sample_rate, buffer_size))
            .collect();

        Self {
            name,
            template_graph,
            voice_instances,
            voices: std::array::from_fn(|_| VoiceState::new()),
            voice_count: DEFAULT_VOICES,
            mix_buffer: vec![0.0; buffer_size * 2], // Stereo
            inputs,
            outputs,
            parameters,
        }
    }

    /// Get mutable reference to template graph (for UI editing)
    pub fn template_graph_mut(&mut self) -> &mut AudioGraph {
        &mut self.template_graph
    }

    /// Get reference to template graph (for serialization)
    pub fn template_graph(&self) -> &AudioGraph {
        &self.template_graph
    }

    /// Rebuild voice instances from template (called after template is edited)
    pub fn rebuild_voices(&mut self) {
        // Clone template to all voice instances
        for voice in &mut self.voice_instances {
            *voice = self.template_graph.clone_graph();

            // Find TemplateInput and TemplateOutput nodes
            let mut template_input_idx = None;
            let mut template_output_idx = None;

            for node_idx in voice.node_indices() {
                if let Some(node) = voice.get_node(node_idx) {
                    match node.node_type() {
                        "TemplateInput" => template_input_idx = Some(node_idx),
                        "TemplateOutput" => template_output_idx = Some(node_idx),
                        _ => {}
                    }
                }
            }

            // Mark ONLY TemplateInput as a MIDI target
            // MIDI will flow through graph connections to other nodes (like MidiToCV)
            if let Some(input_idx) = template_input_idx {
                voice.set_midi_target(input_idx, true);
            }

            // Set TemplateOutput as output node
            voice.set_output_node(template_output_idx);
        }
    }

    /// Find a free voice, or steal one
    /// Priority: inactive → oldest releasing → oldest held
    fn find_voice_for_note_on(&mut self) -> usize {
        // First, look for an inactive voice
        for (i, voice) in self.voices[..self.voice_count].iter().enumerate() {
            if !voice.active {
                return i;
            }
        }

        // No inactive voices — steal the oldest releasing voice
        if let Some((i, _)) = self.voices[..self.voice_count]
            .iter()
            .enumerate()
            .filter(|(_, v)| v.releasing)
            .max_by_key(|(_, v)| v.age)
        {
            return i;
        }

        // No releasing voices either — steal the oldest held voice
        self.voices[..self.voice_count]
            .iter()
            .enumerate()
            .max_by_key(|(_, v)| v.age)
            .map(|(i, _)| i)
            .unwrap_or(0)
    }

    /// Get oscilloscope data from the most relevant voice's subgraph.
    /// Priority: first active voice → first releasing voice → first voice.
    pub fn get_voice_oscilloscope_data(&self, node_id: u32, sample_count: usize) -> Option<(Vec<f32>, Vec<f32>)> {
        let voice_idx = self.best_voice_index();
        let graph = &self.voice_instances[voice_idx];
        let node_idx = petgraph::stable_graph::NodeIndex::new(node_id as usize);
        let audio = graph.get_oscilloscope_data(node_idx, sample_count)?;
        let cv = graph.get_oscilloscope_cv_data(node_idx, sample_count).unwrap_or_default();
        Some((audio, cv))
    }

    /// Find the best voice index to observe: first active → first releasing → 0
    fn best_voice_index(&self) -> usize {
        // First active (non-releasing) voice
        for (i, v) in self.voices[..self.voice_count].iter().enumerate() {
            if v.active && !v.releasing {
                return i;
            }
        }
        // First releasing voice
        for (i, v) in self.voices[..self.voice_count].iter().enumerate() {
            if v.active && v.releasing {
                return i;
            }
        }
        // Fallback to first voice
        0
    }

    /// Find all voices playing a specific note (held, not yet releasing)
    fn find_voices_for_note_off(&self, note: u8) -> Vec<usize> {
        self.voices[..self.voice_count]
            .iter()
            .enumerate()
            .filter_map(|(i, v)| {
                if v.active && !v.releasing && v.note == note {
                    Some(i)
                } else {
                    None
                }
            })
            .collect()
    }
}

impl AudioNode for VoiceAllocatorNode {
    fn category(&self) -> NodeCategory {
        NodeCategory::Utility
    }

    fn inputs(&self) -> &[NodePort] {
        &self.inputs
    }

    fn outputs(&self) -> &[NodePort] {
        &self.outputs
    }

    fn parameters(&self) -> &[Parameter] {
        &self.parameters
    }

    fn set_parameter(&mut self, id: u32, value: f32) {
        match id {
            PARAM_VOICE_COUNT => {
                let new_count = (value.round() as usize).clamp(1, MAX_VOICES);
                if new_count != self.voice_count {
                    self.voice_count = new_count;
                    // Stop voices beyond the new count
                    for voice in &mut self.voices[new_count..] {
                        voice.active = false;
                        voice.releasing = false;
                    }
                }
            }
            _ => {}
        }
    }

    fn get_parameter(&self, id: u32) -> f32 {
        match id {
            PARAM_VOICE_COUNT => self.voice_count as f32,
            _ => 0.0,
        }
    }

    fn handle_midi(&mut self, event: &MidiEvent) {
        let status = event.status & 0xF0;

        match status {
            0x90 => {
                // Note on
                if event.data2 > 0 {
                    let voice_idx = self.find_voice_for_note_on();
                    self.voices[voice_idx].active = true;
                    self.voices[voice_idx].releasing = false;
                    self.voices[voice_idx].note = event.data1;
                    self.voices[voice_idx].note_channel = event.status & 0x0F;
                    self.voices[voice_idx].age = 0;

                    // Store MIDI event for this voice to process
                    self.voices[voice_idx].pending_events.push(*event);
                } else {
                    // Velocity = 0 means note off — mark releasing, keep active for ADSR release
                    let voice_indices = self.find_voices_for_note_off(event.data1);
                    for voice_idx in voice_indices {
                        self.voices[voice_idx].releasing = true;
                        self.voices[voice_idx].pending_events.push(*event);
                    }
                }
            }
            0x80 => {
                // Note off — mark releasing, keep active for ADSR release
                let voice_indices = self.find_voices_for_note_off(event.data1);
                for voice_idx in voice_indices {
                    self.voices[voice_idx].releasing = true;
                    self.voices[voice_idx].pending_events.push(*event);
                }
            }
            _ => {
                // Route to matching-channel voices; channel 0 = global broadcast
                let event_channel = event.status & 0x0F;
                for voice_idx in 0..self.voice_count {
                    let voice = &mut self.voices[voice_idx];
                    if voice.active && (event_channel == 0 || voice.note_channel == event_channel) {
                        voice.pending_events.push(*event);
                    }
                }
            }
        }
    }

    fn process(
        &mut self,
        _inputs: &[&[f32]],
        outputs: &mut [&mut [f32]],
        midi_inputs: &[&[MidiEvent]],
        _midi_outputs: &mut [&mut Vec<MidiEvent>],
        _sample_rate: u32,
    ) {
        // Process MIDI events from input (allocate notes to voices)
        if !midi_inputs.is_empty() {
            for event in midi_inputs[0] {
                self.handle_midi(event);
            }
        }

        if outputs.is_empty() {
            return;
        }

        let output = &mut outputs[0];
        let output_len = output.len();

        // Process each active voice and mix (only up to voice_count)
        for voice_idx in 0..self.voice_count {
            let voice_state = &mut self.voices[voice_idx];
            if voice_state.active {
                voice_state.age = voice_state.age.saturating_add(1);

                // Get pending MIDI events for this voice
                let midi_events = std::mem::take(&mut voice_state.pending_events);

                // IMPORTANT: Process only the slice of mix_buffer that matches output size
                // This prevents phase discontinuities in oscillators
                let mix_slice = &mut self.mix_buffer[..output_len];
                mix_slice.fill(0.0);

                // Process this voice's graph with its MIDI events
                // Note: playback_time is 0.0 since voice allocator doesn't track time
                self.voice_instances[voice_idx].process(mix_slice, &midi_events, crate::time::Beats::ZERO);

                // Auto-deactivate releasing voices that have gone silent
                if voice_state.releasing {
                    let peak = mix_slice.iter().fold(0.0f32, |max, &s| max.max(s.abs()));
                    if peak < 1e-6 {
                        voice_state.active = false;
                        voice_state.releasing = false;
                        continue; // Don't mix silent output
                    }
                }

                // Mix into output (accumulate)
                for (i, sample) in mix_slice.iter().enumerate() {
                    output[i] += sample;
                }
            }
        }
    }

    fn reset(&mut self) {
        for voice in &mut self.voices {
            voice.active = false;
            voice.releasing = false;
            voice.pending_events.clear();
        }
        for graph in &mut self.voice_instances {
            graph.reset();
        }
        self.template_graph.reset();
    }

    fn node_type(&self) -> &str {
        "VoiceAllocator"
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn clone_node(&self) -> Box<dyn AudioNode> {
        // Clone creates a new VoiceAllocator with the same template graph
        // Voice instances will be rebuilt when rebuild_voices() is called
        Box::new(Self {
            name: self.name.clone(),
            template_graph: self.template_graph.clone_graph(),
            voice_instances: self.voice_instances.iter().map(|g| g.clone_graph()).collect(),
            voices: std::array::from_fn(|_| VoiceState::new()), // Reset voices
            voice_count: self.voice_count,
            mix_buffer: vec![0.0; self.mix_buffer.len()],
            inputs: self.inputs.clone(),
            outputs: self.outputs.clone(),
            parameters: self.parameters.clone(),
        })
    }

    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}