//! Strict Standard MIDI File import/export. Unsupported data is reported, never flattened.
use rivet_core::{
    Error, MessageKind, Meter, Note, Performance, PerformanceSettings, Phrase, Result,
    MAX_PLAYBACK_NOTES, PPQ,
};

pub const MAX_FILE_BYTES: usize = 4 * 1024 * 1024;
#[derive(Debug)]
pub struct MidiTrack {
    pub name: String,
    pub phrase: Phrase,
    pub warnings: Vec<String>,
    pub channels: Vec<u8>,
}
#[derive(Debug)]
pub struct MidiFile {
    pub tracks: Vec<MidiTrack>,
    pub tempo: Option<f64>,
    /// The file's time signature; track lengths are rounded to its bars.
    pub meter: Option<Meter>,
}
fn error(s: &str) -> Error {
    Error(s.into())
}

struct Reader<'a> {
    bytes: &'a [u8],
    pos: usize,
}
impl<'a> Reader<'a> {
    fn take(&mut self, len: usize) -> Result<&'a [u8]> {
        let end = self
            .pos
            .checked_add(len)
            .ok_or_else(|| error("MIDI length overflow."))?;
        let out = self
            .bytes
            .get(self.pos..end)
            .ok_or_else(|| error("Truncated MIDI file."))?;
        self.pos = end;
        Ok(out)
    }
    fn byte(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }
    fn u16(&mut self) -> Result<u16> {
        let b = self.take(2)?;
        Ok(u16::from_be_bytes([b[0], b[1]]))
    }
    fn u32(&mut self) -> Result<u32> {
        let b = self.take(4)?;
        Ok(u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
    }
    fn vlq(&mut self) -> Result<u32> {
        let mut v = 0;
        for _ in 0..4 {
            let b = self.byte()?;
            v = (v << 7) | (b & 127) as u32;
            if b & 128 == 0 {
                return Ok(v);
            }
        }
        Err(error("Invalid MIDI variable-length integer."))
    }
    fn data(&mut self) -> Result<u8> {
        let b = self.byte()?;
        if b > 127 {
            Err(error("Invalid MIDI data byte."))
        } else {
            Ok(b)
        }
    }
}
fn ticks(value: u32, division: u16) -> Result<u32> {
    let numerator = value as u64 * PPQ as u64;
    if !numerator.is_multiple_of(division as u64) {
        return Err(error("This MIDI timing cannot be represented exactly at 960 PPQ. Re-export at 96, 480, or 960 PPQ."));
    }
    u32::try_from(numerator / division as u64).map_err(|_| error("MIDI timeline is too long."))
}

pub fn read(bytes: &[u8]) -> Result<MidiFile> {
    if bytes.len() > MAX_FILE_BYTES {
        return Err(error("MIDI file exceeds 4 MiB."));
    }
    let mut r = Reader { bytes, pos: 0 };
    if r.take(4)? != b"MThd" {
        return Err(error("Select a standard .mid MIDI file."));
    }
    let header_len = r.u32()? as usize;
    if header_len < 6 {
        return Err(error("Invalid MIDI header."));
    }
    let format = r.u16()?;
    let count = r.u16()?;
    let division = r.u16()?;
    if format > 1 {
        return Err(error(
            "Asynchronous format-2 MIDI is unsupported. Export format 0 or 1.",
        ));
    }
    if count == 0 || count > 128 || (format == 0 && count != 1) {
        return Err(error("Invalid MIDI track count (maximum 128)."));
    }
    if division == 0 || division & 0x8000 != 0 {
        return Err(error(
            "SMPTE timing is unsupported. Export beat-based MIDI.",
        ));
    }
    r.take(header_len - 6)?;
    let mut tracks = Vec::new();
    let mut tempo = None;
    let mut meter: Option<Meter> = None;
    let mut lengths = Vec::new();
    for track_index in 0..count {
        if r.take(4)? != b"MTrk" {
            return Err(error("Missing MIDI track chunk."));
        }
        let len = r.u32()? as usize;
        let data = r.take(len)?;
        let mut t = Reader {
            bytes: data,
            pos: 0,
        };
        let mut clock = 0u32;
        let mut running = 0u8;
        let mut active = [[None; 128]; 16];
        let mut notes = Vec::new();
        let mut name = format!("Track {}", track_index + 1);
        let mut ignored = false;
        let mut channels = 0u16;
        let mut ended = false;
        let mut events = 0;
        while t.pos < data.len() {
            events += 1;
            if events > 100_000 {
                return Err(error("MIDI track exceeds 100,000 events."));
            }
            clock = clock
                .checked_add(t.vlq()?)
                .ok_or_else(|| error("MIDI timeline overflow."))?;
            let first = t.byte()?;
            let status = if first < 128 {
                if running == 0 {
                    return Err(error("Invalid MIDI running status."));
                }
                t.pos -= 1;
                running
            } else {
                first
            };
            match status {
                0xff => {
                    running = 0;
                    let kind = t.byte()?;
                    let len = t.vlq()? as usize;
                    let body = t.take(len)?;
                    match kind {
                        0x03 => name = String::from_utf8_lossy(body).chars().take(120).collect(),
                        0x2f => {
                            if !body.is_empty() {
                                return Err(error("Invalid end-of-track event."));
                            }
                            ended = true;
                            break;
                        }
                        0x51 => {
                            if body.len() != 3 {
                                return Err(error("Invalid tempo event."));
                            }
                            let us =
                                ((body[0] as u32) << 16) | ((body[1] as u32) << 8) | body[2] as u32;
                            if us == 0 {
                                return Err(error("Invalid zero MIDI tempo."));
                            }
                            if tempo.is_none() {
                                tempo = Some(60_000_000. / us as f64);
                            } else {
                                ignored = true;
                            }
                        }
                        0x58 => {
                            if body.len() != 4 {
                                return Err(error("Invalid time signature."));
                            }
                            let found = Meter {
                                beats: body[0],
                                unit: 1u8.checked_shl(body[1] as u32).unwrap_or(0),
                            };
                            found.validate()?;
                            if meter.is_some_and(|m| m != found) {
                                return Err(error(
                                    "Meter changes inside one file are unsupported. Export one section per meter.",
                                ));
                            }
                            meter = Some(found);
                        }
                        _ => {}
                    }
                }
                0xf0 | 0xf7 => {
                    running = 0;
                    let len = t.vlq()? as usize;
                    t.take(len)?;
                    ignored = true;
                }
                0x80..=0xef => {
                    running = status;
                    let a = t.data()?;
                    let b = if status & 0xf0 == 0xc0 || status & 0xf0 == 0xd0 {
                        0
                    } else {
                        t.data()?
                    };
                    let ch = (status & 15) as usize;
                    match (status & 0xf0, b) {
                        (0x90, 1..=127) => {
                            channels |= 1 << ch;
                            if active[ch][a as usize].is_some() {
                                return Err(error(
                                    "Repeated note-on before note-off is unsupported.",
                                ));
                            }
                            active[ch][a as usize] = Some((clock, b));
                        }
                        (0x80, _) | (0x90, 0) => {
                            let (start, velocity) =
                                active[ch][a as usize].take().ok_or_else(|| {
                                    error("MIDI contains a note-off without a note-on.")
                                })?;
                            let start = ticks(start, division)?;
                            let end = ticks(clock, division)?;
                            if end <= start {
                                return Err(error("MIDI contains a zero-length note."));
                            }
                            if notes.len() >= MAX_PLAYBACK_NOTES {
                                return Err(error("MIDI track has too many notes."));
                            }
                            notes.push(Note {
                                start,
                                duration: end - start,
                                pitch: a,
                                velocity,
                                locked: false,
                                ..Note::default()
                            });
                        }
                        _ => ignored = true,
                    }
                }
                _ => return Err(error("Unsupported MIDI event status.")),
            }
        }
        if !ended {
            return Err(error("MIDI track is missing its end-of-track event."));
        }
        if active.iter().flatten().any(Option::is_some) {
            return Err(error("MIDI contains notes without note-offs."));
        }
        if !notes.is_empty() {
            notes.sort_by_key(|n| (n.start, n.pitch));
            let end = ticks(clock, division)?.max(notes.iter().map(Note::end).max().unwrap_or(0));
            lengths.push(end);
            let warnings = if ignored {
                vec!["Controllers, pitch bend, program changes, SysEx, or later tempo events are not imported. Preview uses a single tempo.".into()]
            } else {
                vec![]
            };
            tracks.push(MidiTrack {
                name,
                phrase: Phrase { notes, length: 0 },
                warnings,
                channels: (0..16).filter(|c| channels & (1 << c) != 0).collect(),
            });
        }
    }
    if tracks.is_empty() {
        return Err(error("This MIDI file has no note tracks."));
    }
    // The time signature may follow the notes in another track, so bars are known only now.
    let bar = meter.unwrap_or_default().bar();
    for (track, end) in tracks.iter_mut().zip(lengths) {
        track.phrase.length = end
            .checked_add(bar - 1)
            .ok_or_else(|| error("MIDI timeline overflow."))?
            / bar
            * bar;
    }
    Ok(MidiFile {
        tracks,
        tempo,
        meter,
    })
}

fn vlq(mut value: u32, out: &mut Vec<u8>) {
    let mut b = [0u8; 5];
    let mut i = 4;
    b[i] = (value & 127) as u8;
    value >>= 7;
    while value > 0 {
        i -= 1;
        b[i] = ((value & 127) as u8) | 128;
        value >>= 7;
    }
    out.extend_from_slice(&b[i..]);
}

fn time_signature(meter: Meter) -> [u8; 8] {
    [
        0,
        0xff,
        0x58,
        4,
        meter.beats,
        meter.unit.trailing_zeros() as u8,
        24,
        8,
    ]
}
fn tempo_event(bpm: f64) -> [u8; 7] {
    let us = (60_000_000. / bpm).round() as u32;
    [
        0,
        0xff,
        0x51,
        3,
        (us >> 16) as u8,
        (us >> 8) as u8,
        us as u8,
    ]
}

pub fn write(phrase: &Phrase, bpm: f64, meter: Meter) -> Result<Vec<u8>> {
    phrase.validate(MAX_PLAYBACK_NOTES)?;
    meter.validate()?;
    if !bpm.is_finite() || !(30.0..=300.0).contains(&bpm) {
        return Err(error("Tempo must be 30–300 BPM."));
    }
    let mut track = tempo_event(bpm).to_vec();
    track.extend_from_slice(&time_signature(meter));
    let mut events = Vec::new();
    for n in &phrase.notes {
        events.push((n.start, 1, n.pitch, n.velocity));
        events.push((n.end(), 0, n.pitch, 0));
    }
    events.sort_unstable();
    let mut previous = 0;
    for (time, on, pitch, velocity) in events {
        vlq(time - previous, &mut track);
        track.extend_from_slice(&[if on == 1 { 0x90 } else { 0x80 }, pitch, velocity]);
        previous = time;
    }
    vlq(phrase.length - previous, &mut track);
    track.extend_from_slice(&[0xff, 0x2f, 0]);
    let mut out = b"MThd\0\0\0\x06\0\0\0\x01".to_vec();
    out.extend_from_slice(&(PPQ as u16).to_be_bytes());
    out.extend_from_slice(b"MTrk");
    out.extend_from_slice(&(track.len() as u32).to_be_bytes());
    out.extend_from_slice(&track);
    Ok(out)
}

/// Two explicit sound roles, named tracks, channel routing, and mapped expression.
/// Compensated export adds a full bar before the intended musical start.
pub fn write_performance(
    phrase: &Phrase,
    settings: &PerformanceSettings,
    bpm: f64,
    meter: Meter,
) -> Result<Vec<u8>> {
    meter.validate()?;
    let p = Performance::compile(phrase, settings, bpm, settings.compensated_export, meter)?;
    let shift = if settings.compensated_export && p.lead_in > 0 {
        meter.bar()
    } else {
        0
    };
    let mut conductor = tempo_event(bpm).to_vec();
    conductor.extend_from_slice(&time_signature(meter));
    vlq(phrase.length + shift, &mut conductor);
    conductor.extend_from_slice(&[0xff, 0x2f, 0]);
    let mut tracks = vec![conductor];
    for role in 0..2 {
        let events: Vec<_> = p
            .events
            .iter()
            .filter(|e| e.role as usize == role && !e.next_cycle)
            .collect();
        if events.is_empty() {
            continue;
        }
        let mut track = vec![0, 0xff, 3];
        let name = format!(
            "{} (channel {})",
            settings.voices[role].name,
            settings.voices[role].channel + 1
        );
        vlq(name.len() as u32, &mut track);
        track.extend_from_slice(name.as_bytes());
        let mut previous = 0;
        for e in events {
            let tick = (e.tick + shift as i32) as u32;
            vlq(tick - previous, &mut track);
            let kind = match e.message.kind {
                MessageKind::Off => 0x80,
                MessageKind::On => 0x90,
                MessageKind::Control => 0xb0,
            };
            track.extend_from_slice(&[kind | e.message.channel, e.message.key, e.message.value]);
            previous = tick;
        }
        let end = phrase.length + shift;
        if settings.voices[role].controller.enabled {
            vlq(end - previous, &mut track);
            track.extend_from_slice(&[
                0xb0 | settings.voices[role].channel,
                settings.voices[role].controller.cc,
                settings.voices[role].controller.rest,
            ]);
            previous = end;
        }
        vlq(end - previous, &mut track);
        track.extend_from_slice(&[0xff, 0x2f, 0]);
        tracks.push(track);
    }
    let mut out = b"MThd\0\0\0\x06\0\x01".to_vec();
    out.extend_from_slice(&(tracks.len() as u16).to_be_bytes());
    out.extend_from_slice(&(PPQ as u16).to_be_bytes());
    for track in tracks {
        out.extend_from_slice(b"MTrk");
        out.extend_from_slice(&(track.len() as u32).to_be_bytes());
        out.extend_from_slice(&track);
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    fn file(track: &[u8]) -> Vec<u8> {
        let mut b = b"MThd\0\0\0\x06\0\0\0\x01\x03\xc0MTrk".to_vec();
        b.extend_from_slice(&(track.len() as u32).to_be_bytes());
        b.extend_from_slice(track);
        b
    }
    #[test]
    fn round_trip_source_and_generated_phrase() {
        let source = rivet_core::demo_motif();
        let bytes = write(&source, 140., Meter::default()).unwrap();
        let parsed = read(&bytes).unwrap();
        assert_eq!(parsed.tracks[0].phrase, source);
        assert!((parsed.tempo.unwrap() - 140.).abs() < 0.001);
        assert_eq!(parsed.meter, Some(Meter::default()));
        let story = rivet_core::Session::demo().compose().unwrap();
        assert_eq!(
            read(&write(&story, 140., Meter::default()).unwrap())
                .unwrap()
                .tracks[0]
                .phrase,
            story
        );
    }
    #[test]
    fn running_status_and_zero_velocity_off() {
        let parsed = read(&file(&[
            0, 0x90, 60, 90, 0x83, 0x60, 60, 0, 0, 0xff, 0x2f, 0,
        ]))
        .unwrap();
        assert_eq!(parsed.tracks[0].phrase.notes[0].duration, 480);
    }
    #[test]
    fn malformed_and_unsupported_files() {
        let good = write(&rivet_core::demo_motif(), 140., Meter::default()).unwrap();
        for n in 0..good.len() {
            assert!(read(&good[..n]).is_err());
        }
        let mut smpte = good.clone();
        smpte[12] = 0x80;
        assert!(read(&smpte).unwrap_err().0.contains("SMPTE"));
        assert!(read(&file(&[0, 0x90, 60, 80, 0, 0xff, 0x2f, 0])).is_err());
        assert!(
            read(&file(&[0, 0xff, 0x58, 4, 11, 3, 24, 8, 0, 0xff, 0x2f, 0]))
                .unwrap_err()
                .0
                .contains("Supported meters")
        );
        assert!(read(&file(&[0x81, 0x81, 0x81, 0x81, 0, 0xff, 0x2f, 0])).is_err());
    }
    #[test]
    fn overlap_is_available_for_selection_but_rejected_as_motif() {
        let p = read(&file(&[
            0, 0x90, 60, 90, 0, 0x90, 67, 80, 10, 0x80, 60, 0, 0, 0x80, 67, 0, 0, 0xff, 0x2f, 0,
        ]))
        .unwrap();
        assert!(p.tracks[0]
            .phrase
            .validate(rivet_core::MAX_SOURCE_NOTES)
            .unwrap_err()
            .0
            .contains("Overlapping"));
    }
    #[test]
    fn separate_tracks_and_timing_resolution() {
        let a = write(&rivet_core::demo_motif(), 140., Meter::default()).unwrap();
        let mut two = a[..14].to_vec();
        two[9] = 1;
        two[11] = 2;
        two.extend_from_slice(&a[14..]);
        two.extend_from_slice(&a[14..]);
        assert_eq!(read(&two).unwrap().tracks.len(), 2);
        let mut b = file(&[1, 0x90, 60, 80, 1, 0x80, 60, 0, 0, 0xff, 0x2f, 0]);
        b[12..14].copy_from_slice(&1000u16.to_be_bytes());
        assert!(read(&b).unwrap_err().0.contains("exactly"));
    }

    #[test]
    fn performance_export_has_named_channel_tracks_and_controller_resets() {
        let mut settings = PerformanceSettings {
            mode: rivet_core::RoleMode::Alternating,
            ..Default::default()
        };
        for (role, v) in settings.voices.iter_mut().enumerate() {
            v.controller.enabled = true;
            v.controller.cc = 20 + role as u8;
            v.controller.rest = 5 + role as u8;
        }
        let source = rivet_core::demo_motif();
        let bytes = write_performance(&source, &settings, 120., Meter::default()).unwrap();
        let file = read(&bytes).unwrap();
        assert_eq!(file.tracks.len(), 2);
        assert_eq!(file.tracks[0].channels, vec![0]);
        assert_eq!(file.tracks[1].channels, vec![1]);
        assert!(file.tracks[0].name.contains("Primary"));
        assert!(file.tracks[1].name.contains("Answer"));
        assert_eq!(
            file.tracks
                .iter()
                .map(|t| t.phrase.notes.len())
                .sum::<usize>(),
            source.notes.len()
        );
        assert!(file.tracks.iter().all(|t| !t.warnings.is_empty()));
        assert!(bytes
            .windows(7)
            .any(|p| p == [0xb0, 20, 5, 0, 0xff, 0x2f, 0]));
        assert!(bytes
            .windows(7)
            .any(|p| p == [0xb1, 21, 6, 0, 0xff, 0x2f, 0]));
    }

    #[test]
    fn compensated_export_adds_lead_in_while_default_export_keeps_the_grid() {
        let source = rivet_core::demo_motif();
        let mut settings = PerformanceSettings::default();
        settings.voices[0].attack_ms = 100;
        let direct =
            read(&write_performance(&source, &settings, 120., Meter::default()).unwrap()).unwrap();
        assert_eq!(direct.tracks[0].phrase.notes[0].start, 0);
        settings.compensated_export = true;
        let compensated =
            read(&write_performance(&source, &settings, 120., Meter::default()).unwrap()).unwrap();
        assert_eq!(
            compensated.tracks[0].phrase.notes[0].start,
            rivet_core::BAR - 192
        );
        assert_eq!(
            compensated.tracks[0].phrase.length,
            source.length + rivet_core::BAR
        );
    }
    #[test]
    fn compound_meter_round_trips_and_sets_bar_length() {
        let meter = Meter { beats: 6, unit: 8 };
        let phrase = Phrase {
            length: 2 * meter.bar(),
            notes: vec![Note {
                start: meter.beat(),
                duration: 480,
                pitch: 62,
                velocity: 80,
                ..Default::default()
            }],
        };
        let parsed = read(&write(&phrase, 120., meter).unwrap()).unwrap();
        assert_eq!(parsed.meter, Some(meter));
        assert_eq!(parsed.tracks[0].phrase, phrase);
        let three = read(&file(&[
            0, 0xff, 0x58, 4, 3, 2, 24, 8, 0, 0x90, 60, 80, 0x83, 0x60, 0x80, 60, 0, 0, 0xff, 0x2f,
            0,
        ]))
        .unwrap();
        assert_eq!(three.meter, Some(Meter { beats: 3, unit: 4 }));
        assert_eq!(three.tracks[0].phrase.length, 3 * PPQ);
    }
}