libfreemkv 0.7.2

//! MPLS playlist parser — Blu-ray movie playlists.
//!
//! Each .mpls file in BDMV/PLAYLIST/ defines a title.
//! Contains play items (clips) with in/out timestamps,
//! stream info (video, audio, subtitle tracks).
//!
//! Reference: https://github.com/lw/BluRay/wiki/MPLS

use crate::error::{Error, Result};

/// Parsed MPLS playlist.
#[derive(Debug)]
pub struct Playlist {
    /// MPLS version (e.g. "0200" or "0300")
    pub version: String,
    /// Play items in playback order
    pub play_items: Vec<PlayItem>,
    /// Streams from the first play item's STN table
    pub streams: Vec<StreamEntry>,
}

/// A play item — one clip reference with in/out times.
#[derive(Debug)]
pub struct PlayItem {
    /// Clip filename without extension (e.g. "00001")
    pub clip_id: String,
    /// In-time in 45kHz ticks
    pub in_time: u32,
    /// Out-time in 45kHz ticks
    pub out_time: u32,
    /// Connection condition (1=seamless, 5/6=non-seamless)
    pub connection_condition: u8,
}

/// A stream entry from the STN table.
#[derive(Debug, Clone)]
pub struct StreamEntry {
    /// Stream category: 1=video, 2=audio, 3=PG subtitle, 4=IG, 5=secondary audio, 6=secondary video, 7=DV EL
    pub stream_type: u8,
    /// MPEG-TS PID
    pub pid: u16,
    /// Coding type (0x24=HEVC, 0x1B=H264, 0x83=TrueHD, etc.)
    pub coding_type: u8,
    /// Video format (1=480i, 4=1080i, 5=720p, 6=1080p, 8=2160p)
    pub video_format: u8,
    /// Video frame rate (1=23.976, 2=24, 3=25, 4=29.97, 6=50, 7=59.94)
    pub video_rate: u8,
    /// Audio channel layout (1=mono, 3=stereo, 6=5.1, 12=7.1)
    pub audio_format: u8,
    /// Audio sample rate (1=48kHz, 4=96kHz, 5=192kHz)
    pub audio_rate: u8,
    /// ISO 639-2 language code (e.g. "eng")
    pub language: String,
    /// HDR dynamic range (0=SDR, 1=HDR10, 2=Dolby Vision)
    pub dynamic_range: u8,
    /// Color space (0=unknown, 1=BT.709, 2=BT.2020)
    pub color_space: u8,
    /// Whether this is a secondary stream (commentary, PiP, DV EL)
    pub secondary: bool,
}

/// Parse an MPLS file from raw bytes.
pub fn parse(data: &[u8]) -> Result<Playlist> {
    if data.len() < 40 {
        return Err(Error::MplsParse);
    }
    if &data[0..4] != b"MPLS" {
        return Err(Error::MplsParse);
    }

    let version = String::from_utf8_lossy(&data[4..8]).to_string();
    let playlist_start = u32::from_be_bytes([data[8], data[9], data[10], data[11]]) as usize;

    if playlist_start + 10 > data.len() {
        return Err(Error::MplsParse);
    }

    let pl = &data[playlist_start..];
    let num_play_items = u16::from_be_bytes([pl[6], pl[7]]) as usize;

    let mut play_items = Vec::with_capacity(num_play_items);
    let mut streams = Vec::new();
    let mut pos = 10;

    for item_idx in 0..num_play_items {
        if pos + 2 > pl.len() { break; }
        let item_length = u16::from_be_bytes([pl[pos], pl[pos + 1]]) as usize;
        if pos + 2 + item_length > pl.len() { break; }

        let item = &pl[pos + 2..pos + 2 + item_length];
        if item.len() < 20 { pos += 2 + item_length; continue; }

        let clip_id = String::from_utf8_lossy(&item[0..5]).to_string();
        let connection_condition = item[9] & 0x0F;
        let in_time = u32::from_be_bytes([item[12], item[13], item[14], item[15]]);
        let out_time = u32::from_be_bytes([item[16], item[17], item[18], item[19]]);

        // Parse STN table from the first play item
        // PlayItem layout after out_time:
        //   [20:28] UO_mask_table (8 bytes)
        //   [28]    misc flags (1 byte)
        //   [29]    still_mode (1 byte)
        //   [30:32] still_time (2 bytes)
        //   [32:]   STN_table
        const STN_OFFSET: usize = 32;
        if item_idx == 0 && item.len() > STN_OFFSET + 16 {
            // STN header: length(2) + reserved(2) + counts(8) + reserved(4) = 16 bytes
            let n_video = item[STN_OFFSET + 4] as usize;
            let n_audio = item[STN_OFFSET + 5] as usize;
            let n_pg = item[STN_OFFSET + 6] as usize;
            let n_ig = item[STN_OFFSET + 7] as usize;
            let n_sec_audio = item[STN_OFFSET + 8] as usize;
            let n_sec_video = item[STN_OFFSET + 9] as usize;
            let n_pip_pg = item[STN_OFFSET + 10] as usize;
            let n_dv = item[STN_OFFSET + 11] as usize;

            let mut spos = STN_OFFSET + 16;

            // Primary video
            for _ in 0..n_video {
                if let Some((entry, next)) = parse_stream_entry(item, spos, 1) {
                    streams.push(entry);
                    spos = next;
                } else { break; }
            }
            // Primary audio
            for _ in 0..n_audio {
                if let Some((entry, next)) = parse_stream_entry(item, spos, 2) {
                    streams.push(entry);
                    spos = next;
                } else { break; }
            }
            // PG subtitles
            for _ in 0..n_pg {
                if let Some((entry, next)) = parse_stream_entry(item, spos, 3) {
                    streams.push(entry);
                    spos = next;
                } else { break; }
            }
            // IG (skip but advance)
            for _ in 0..n_ig {
                if let Some((_, next)) = parse_stream_entry(item, spos, 4) {
                    spos = next;
                } else { break; }
            }
            // Secondary audio
            for _ in 0..n_sec_audio {
                if let Some((mut entry, next)) = parse_stream_entry(item, spos, 2) {
                    entry.stream_type = 5;
                    entry.secondary = true;
                    streams.push(entry);
                    // Skip extra ref bytes: num_refs(1) + reserved(1) + refs + padding
                    if next < item.len() {
                        let n_refs = item[next] as usize;
                        spos = next + 2 + n_refs + (n_refs % 2);
                    } else { spos = next; }
                } else { break; }
            }
            // Secondary video (PiP)
            for _ in 0..n_sec_video {
                if let Some((mut entry, next)) = parse_stream_entry(item, spos, 1) {
                    entry.stream_type = 6;
                    entry.secondary = true;
                    streams.push(entry);
                    // Skip extra ref bytes (audio refs + PG refs)
                    if next + 2 < item.len() {
                        let n_arefs = item[next] as usize;
                        let after_arefs = next + 2 + n_arefs + (n_arefs % 2);
                        if after_arefs < item.len() {
                            let n_prefs = item[after_arefs] as usize;
                            spos = after_arefs + 2 + n_prefs + (n_prefs % 2);
                        } else { spos = after_arefs; }
                    } else { spos = next; }
                } else { break; }
            }
            // Secondary PG (PiP subtitles) — must consume to keep spos aligned
            for _ in 0..n_pip_pg {
                if let Some((mut entry, next)) = parse_stream_entry(item, spos, 3) {
                    entry.secondary = true;
                    streams.push(entry);
                    // Skip reference data: num_refs(1) + reserved(1) + refs + padding
                    if next < item.len() {
                        let n_refs = item[next] as usize;
                        spos = next + 2 + n_refs + (n_refs % 2);
                    } else { spos = next; }
                } else { break; }
            }
            // Dolby Vision enhancement layer
            for _ in 0..n_dv {
                if let Some((mut entry, next)) = parse_stream_entry(item, spos, 1) {
                    entry.stream_type = 7;
                    entry.secondary = true;
                    streams.push(entry);
                    spos = next;
                } else { break; }
            }
        }

        play_items.push(PlayItem {
            clip_id,
            in_time,
            out_time,
            connection_condition,
        });

        pos += 2 + item_length;
    }

    Ok(Playlist {
        version,
        play_items,
        streams,
    })
}

/// Parse one stream entry from the STN table.
/// Returns (StreamEntry, next position) or None.
fn parse_stream_entry(item: &[u8], pos: usize, stream_type: u8) -> Option<(StreamEntry, usize)> {
    if pos + 2 > item.len() { return None; }

    // Stream entry: length(1) + data
    let se_len = item[pos] as usize;
    let se_end = pos + 1 + se_len;
    if se_end > item.len() { return None; }

    // PID from stream entry (type 0x01 = PlayItem stream: PID at bytes 2-3)
    let pid = if item[pos + 1] == 0x01 && pos + 4 <= item.len() {
        u16::from_be_bytes([item[pos + 2], item[pos + 3]])
    } else {
        0
    };

    // Stream attributes: length(1) + coding_type(1) + format-specific data
    if se_end + 2 > item.len() { return None; }
    let sa_len = item[se_end] as usize;
    let sa_end = se_end + 1 + sa_len;
    if sa_end > item.len() || sa_len < 1 { return None; }

    let sa = &item[se_end + 1..se_end + 1 + sa_len];
    let coding_type = sa[0];


    let mut video_format = 0u8;
    let mut video_rate = 0u8;
    let mut audio_format = 0u8;
    let mut audio_rate = 0u8;
    let mut dynamic_range = 0u8;
    let mut color_space_val = 0u8;
    let mut language = String::new();

    match stream_type {
        1 => {
            // Video: coding_type(1) + format_rate(1) + [hdr_info(1) if HEVC]
            if sa.len() >= 2 {
                video_format = (sa[1] >> 4) & 0x0F;
                video_rate = sa[1] & 0x0F;
            }
            if coding_type == 0x24 && sa.len() > 2 {
                dynamic_range = (sa[2] >> 4) & 0x0F;
                color_space_val = sa[2] & 0x0F;
            }
        }
        2 => {
            // Audio: coding_type(1) + format_rate(1) + language(3)
            // Exception: PGS (0x90/0x91) in audio slot uses PG layout: coding_type(1) + language(3)
            if coding_type == 0x90 || coding_type == 0x91 {
                if sa.len() >= 4 {
                    language = String::from_utf8_lossy(&sa[1..4]).to_string();
                }
            } else {
                if sa.len() >= 2 {
                    audio_format = (sa[1] >> 4) & 0x0F;
                    audio_rate = sa[1] & 0x0F;
                }
                if sa.len() >= 5 {
                    language = String::from_utf8_lossy(&sa[2..5]).to_string();
                }
            }
        }
        3 | 4 => {
            // PG/IG: coding_type(1) + language(3)
            if sa.len() >= 4 {
                language = String::from_utf8_lossy(&sa[1..4]).to_string();
            }
        }
        5 => {
            // Secondary audio: same as primary audio
            if sa.len() >= 2 {
                audio_format = (sa[1] >> 4) & 0x0F;
                audio_rate = sa[1] & 0x0F;
            }
            if sa.len() >= 5 {
                language = String::from_utf8_lossy(&sa[2..5]).to_string();
            }
        }
        6 | 7 => {
            // Secondary video: same as primary video
            if sa.len() >= 2 {
                video_format = (sa[1] >> 4) & 0x0F;
                video_rate = sa[1] & 0x0F;
            }
            if coding_type == 0x24 && sa.len() > 2 {
                dynamic_range = (sa[2] >> 4) & 0x0F;
                color_space_val = sa[2] & 0x0F;
            }
        }
        _ => {}
    }

    Some((StreamEntry {
        stream_type,
        pid,
        coding_type,
        video_format,
        video_rate,
        audio_format,
        audio_rate,
        language,
        dynamic_range,
        color_space: color_space_val,
        secondary: false,
    }, sa_end))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Build a minimal MPLS binary with given play items and STN streams on the first item.
    /// STN counts: (n_video, n_audio, n_pg, n_ig, n_sec_audio, n_sec_video, n_pip_pg, n_dv)
    fn build_mpls(
        play_items_data: &[(/*clip_id*/&[u8;5], /*conn*/u8, /*in_time*/u32, /*out_time*/u32)],
        stn_counts: (u8, u8, u8, u8, u8, u8, u8, u8),
        stream_entries: &[Vec<u8>], // raw stream entry + attributes bytes for each stream
    ) -> Vec<u8> {
        let playlist_start: u32 = 40; // right after the 40-byte header
        let mut buf = Vec::new();

        // File header: "MPLS" + version + playlist_start + 3 more offsets (unused)
        buf.extend_from_slice(b"MPLS0200");
        buf.extend_from_slice(&playlist_start.to_be_bytes());
        // mark_start, extension_start (unused by parser), padding to 40 bytes
        buf.extend_from_slice(&[0u8; 28]);

        // PlayList section starts here (offset 40)
        // PlayList: length(4) + reserved(2) + num_play_items(2) + num_sub_paths(2) = 10 header bytes
        let pl_start = buf.len();
        buf.extend_from_slice(&[0u8; 4]); // length placeholder
        buf.extend_from_slice(&[0u8; 2]); // reserved
        buf.extend_from_slice(&(play_items_data.len() as u16).to_be_bytes());
        buf.extend_from_slice(&[0u8; 2]); // num_sub_paths

        for (idx, (clip_id, conn, in_time, out_time)) in play_items_data.iter().enumerate() {
            // Build play item content
            let mut item = Vec::new();
            // [0..5] clip_id
            item.extend_from_slice(*clip_id);
            // [5..9] codec_id ("M2TS")
            item.extend_from_slice(b"M2TS");
            // [9] connection_condition in low nibble
            item.push(*conn & 0x0F);
            // [10..12] reserved
            item.extend_from_slice(&[0u8; 2]);
            // [12..16] in_time
            item.extend_from_slice(&in_time.to_be_bytes());
            // [16..20] out_time
            item.extend_from_slice(&out_time.to_be_bytes());
            // [20..28] UO_mask_table
            item.extend_from_slice(&[0u8; 8]);
            // [28] misc flags
            item.push(0);
            // [29] still_mode
            item.push(0);
            // [30..32] still_time
            item.extend_from_slice(&[0u8; 2]);

            // STN table (only for the first play item)
            if idx == 0 {
                // STN header: length(2) + reserved(2) + counts(8) + reserved(4) = 16 bytes
                let stn_header_start = item.len();
                item.extend_from_slice(&[0u8; 2]); // STN length placeholder
                item.extend_from_slice(&[0u8; 2]); // reserved
                item.push(stn_counts.0); // n_video
                item.push(stn_counts.1); // n_audio
                item.push(stn_counts.2); // n_pg
                item.push(stn_counts.3); // n_ig
                item.push(stn_counts.4); // n_sec_audio
                item.push(stn_counts.5); // n_sec_video
                item.push(stn_counts.6); // n_pip_pg
                item.push(stn_counts.7); // n_dv
                item.extend_from_slice(&[0u8; 4]); // reserved

                // Stream entries
                for se in stream_entries {
                    item.extend_from_slice(se);
                }

                // Patch STN length
                let stn_len = (item.len() - stn_header_start - 2) as u16;
                let stn_len_bytes = stn_len.to_be_bytes();
                item[stn_header_start] = stn_len_bytes[0];
                item[stn_header_start + 1] = stn_len_bytes[1];
            }

            // Write item_length(2) + item
            let item_length = item.len() as u16;
            buf.extend_from_slice(&item_length.to_be_bytes());
            buf.extend_from_slice(&item);
        }

        // Patch PlayList length
        let pl_len = (buf.len() - pl_start - 4) as u32;
        let pl_len_bytes = pl_len.to_be_bytes();
        buf[pl_start] = pl_len_bytes[0];
        buf[pl_start + 1] = pl_len_bytes[1];
        buf[pl_start + 2] = pl_len_bytes[2];
        buf[pl_start + 3] = pl_len_bytes[3];

        buf
    }

    /// Build a stream entry (stream_entry part + stream_attributes part).
    /// stream_entry: type=0x01 (PlayItem stream), PID given.
    /// For video: attrs = coding_type(1) + format_rate(1) [+ hdr_byte if HEVC]
    /// For audio: attrs = coding_type(1) + format_rate(1) + language(3)
    /// For PG:    attrs = coding_type(1) + language(3)
    fn build_stream_entry_video(pid: u16, coding_type: u8, format: u8, rate: u8, hdr: Option<u8>) -> Vec<u8> {
        let mut out = Vec::new();
        // Stream entry: length(1) + sub_path_type(1) + pid(2)
        out.push(3); // se_len = 3 bytes (type + pid_hi + pid_lo)
        out.push(0x01); // type: PlayItem stream
        out.extend_from_slice(&pid.to_be_bytes());
        // Stream attributes
        let mut attrs = vec![coding_type, (format << 4) | rate];
        if let Some(h) = hdr {
            attrs.push(h);
        }
        out.push(attrs.len() as u8); // sa_len
        out.extend_from_slice(&attrs);
        out
    }

    fn build_stream_entry_audio(pid: u16, coding_type: u8, ch_layout: u8, sample_rate: u8, lang: &[u8; 3]) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(3);
        out.push(0x01);
        out.extend_from_slice(&pid.to_be_bytes());
        // attrs: coding_type(1) + format_rate(1) + language(3)
        let attrs = vec![coding_type, (ch_layout << 4) | sample_rate, lang[0], lang[1], lang[2]];
        out.push(attrs.len() as u8);
        out.extend_from_slice(&attrs);
        out
    }

    fn build_stream_entry_pg(pid: u16, coding_type: u8, lang: &[u8; 3]) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(3);
        out.push(0x01);
        out.extend_from_slice(&pid.to_be_bytes());
        // attrs: coding_type(1) + language(3)
        let attrs = vec![coding_type, lang[0], lang[1], lang[2]];
        out.push(attrs.len() as u8);
        out.extend_from_slice(&attrs);
        out
    }

    #[test]
    fn parse_valid_mpls() {
        let in_time: u32 = 90000;   // 2 seconds at 45kHz
        let out_time: u32 = 4500000; // 100 seconds

        let video = build_stream_entry_video(0x1011, 0x1B, 6, 1, None); // H264, 1080p, 23.976
        let audio = build_stream_entry_audio(0x1100, 0x83, 6, 1, b"eng"); // TrueHD, 5.1, 48kHz
        let pg = build_stream_entry_pg(0x1200, 0x90, b"eng"); // PGS subtitle

        let data = build_mpls(
            &[(b"00001", 1, in_time, out_time)],
            (1, 1, 1, 0, 0, 0, 0, 0),
            &[video, audio, pg],
        );

        let playlist = parse(&data).expect("should parse valid MPLS");
        assert_eq!(playlist.version, "0200");
        assert_eq!(playlist.play_items.len(), 1);
        assert_eq!(playlist.play_items[0].clip_id, "00001");
        assert_eq!(playlist.play_items[0].in_time, in_time);
        assert_eq!(playlist.play_items[0].out_time, out_time);
        assert_eq!(playlist.play_items[0].connection_condition, 1);
    }

    #[test]
    fn parse_streams() {
        let video = build_stream_entry_video(0x1011, 0x24, 8, 1, Some(0x12)); // HEVC, 2160p, 23.976, HDR10+BT.2020
        let audio = build_stream_entry_audio(0x1100, 0x83, 6, 1, b"eng");
        let pg = build_stream_entry_pg(0x1200, 0x90, b"fra");

        let data = build_mpls(
            &[(b"00001", 1, 0, 9000000)],
            (1, 1, 1, 0, 0, 0, 0, 0),
            &[video, audio, pg],
        );

        let playlist = parse(&data).expect("should parse");
        assert_eq!(playlist.streams.len(), 3);

        // Video stream
        let v = &playlist.streams[0];
        assert_eq!(v.stream_type, 1);
        assert_eq!(v.pid, 0x1011);
        assert_eq!(v.coding_type, 0x24); // HEVC
        assert_eq!(v.video_format, 8);   // 2160p
        assert_eq!(v.video_rate, 1);     // 23.976
        assert_eq!(v.dynamic_range, 1);  // HDR10
        assert_eq!(v.color_space, 2);    // BT.2020
        assert!(!v.secondary);

        // Audio stream
        let a = &playlist.streams[1];
        assert_eq!(a.stream_type, 2);
        assert_eq!(a.pid, 0x1100);
        assert_eq!(a.coding_type, 0x83); // TrueHD
        assert_eq!(a.audio_format, 6);   // 5.1
        assert_eq!(a.audio_rate, 1);     // 48kHz
        assert_eq!(a.language, "eng");
        assert!(!a.secondary);

        // PG subtitle stream
        let s = &playlist.streams[2];
        assert_eq!(s.stream_type, 3);
        assert_eq!(s.pid, 0x1200);
        assert_eq!(s.coding_type, 0x90); // PGS
        assert_eq!(s.language, "fra");
        assert!(!s.secondary);
    }

    #[test]
    fn parse_invalid_magic() {
        let mut data = build_mpls(
            &[(b"00001", 1, 0, 9000000)],
            (0, 0, 0, 0, 0, 0, 0, 0),
            &[],
        );
        data[0] = b'X';
        data[1] = b'X';
        data[2] = b'X';
        data[3] = b'X';
        assert!(parse(&data).is_err());
    }

    #[test]
    fn parse_truncated() {
        // Less than 40 bytes
        assert!(parse(&[0u8; 10]).is_err());
        assert!(parse(b"MPLS0200").is_err());
        assert!(parse(&[0u8; 39]).is_err());
    }

    #[test]
    fn parse_multiple_play_items() {
        let video = build_stream_entry_video(0x1011, 0x1B, 6, 1, None);

        let data = build_mpls(
            &[
                (b"00001", 1, 90000, 4500000),
                (b"00002", 5, 4500000, 9000000),
                (b"00003", 6, 9000000, 13500000),
            ],
            (1, 0, 0, 0, 0, 0, 0, 0),
            &[video],
        );

        let playlist = parse(&data).expect("should parse multiple play items");
        assert_eq!(playlist.play_items.len(), 3);
        assert_eq!(playlist.play_items[0].clip_id, "00001");
        assert_eq!(playlist.play_items[0].connection_condition, 1);
        assert_eq!(playlist.play_items[1].clip_id, "00002");
        assert_eq!(playlist.play_items[1].connection_condition, 5);
        assert_eq!(playlist.play_items[1].in_time, 4500000);
        assert_eq!(playlist.play_items[2].clip_id, "00003");
        assert_eq!(playlist.play_items[2].connection_condition, 6);
        assert_eq!(playlist.play_items[2].out_time, 13500000);
    }

    #[test]
    fn parse_secondary_streams() {
        // Primary video
        let video = build_stream_entry_video(0x1011, 0x1B, 6, 1, None);
        // Secondary audio (stream_type 5): build as audio, parser overrides type to 5
        let sec_audio_se = build_stream_entry_audio(0x1A00, 0x83, 3, 1, b"eng");
        // Need ref bytes after secondary audio: num_refs(1) + reserved(1) = 2 bytes min
        let mut sec_audio_with_refs = sec_audio_se;
        sec_audio_with_refs.push(0); // num_refs = 0
        sec_audio_with_refs.push(0); // reserved

        // Secondary video (stream_type 6): build as video, parser overrides type to 6
        let sec_video_se = build_stream_entry_video(0x1B00, 0x1B, 4, 1, None);
        // Need ref bytes: n_arefs(1) + reserved(1) + n_prefs(1) + reserved(1) = 4 bytes
        let mut sec_video_with_refs = sec_video_se;
        sec_video_with_refs.push(0); // n_arefs = 0
        sec_video_with_refs.push(0); // reserved
        sec_video_with_refs.push(0); // n_prefs = 0
        sec_video_with_refs.push(0); // reserved

        let data = build_mpls(
            &[(b"00001", 1, 0, 9000000)],
            (1, 0, 0, 0, 1, 1, 0, 0), // 1 video, 0 audio, 0 pg, 0 ig, 1 sec_audio, 1 sec_video
            &[video, sec_audio_with_refs, sec_video_with_refs],
        );

        let playlist = parse(&data).expect("should parse secondary streams");
        // Should have 3 streams: primary video, secondary audio, secondary video
        assert_eq!(playlist.streams.len(), 3);

        // Primary video
        assert_eq!(playlist.streams[0].stream_type, 1);
        assert!(!playlist.streams[0].secondary);

        // Secondary audio
        assert_eq!(playlist.streams[1].stream_type, 5);
        assert!(playlist.streams[1].secondary);
        assert_eq!(playlist.streams[1].pid, 0x1A00);

        // Secondary video
        assert_eq!(playlist.streams[2].stream_type, 6);
        assert!(playlist.streams[2].secondary);
        assert_eq!(playlist.streams[2].pid, 0x1B00);
    }
}