libbitsub-core 1.10.1

//! VobSub parser.

use memchr::memchr;
use std::collections::HashMap;

use super::{
    DebandConfig, ExtractedVobSub, IdxParseResult, SubtitlePacket, VobSubPalette, VobSubTimestamp,
    apply_deband, decode_vobsub_rle, extract_vobsub_from_mks, parse_idx, parse_subtitle_packet,
};
use crate::utils::binary_search_timestamp;

/// VobSub subtitle parser and renderer.
pub struct VobSubParser {
    /// Parsed IDX data
    idx_data: Option<IdxParseResult>,
    /// Raw SUB file data
    sub_data: Option<Vec<u8>>,
    /// Timestamps in milliseconds for quick lookup
    timestamps_ms: Vec<u32>,
    /// Cache for decoded subtitle packets
    packet_cache: HashMap<usize, Option<SubtitlePacket>>,
    /// Debanding configuration
    deband_config: DebandConfig,
    /// Whether the parser was loaded from IDX metadata.
    loaded_from_idx: bool,
    /// Last non-fatal render issue for diagnostics.
    last_render_issue: Option<String>,
}

impl VobSubParser {
    /// Create a new VobSub parser.
    pub fn new() -> Self {
        Self {
            idx_data: None,
            sub_data: None,
            timestamps_ms: Vec::new(),
            packet_cache: HashMap::new(),
            deband_config: DebandConfig::default(),
            loaded_from_idx: false,
            last_render_issue: None,
        }
    }

    /// Load VobSub from IDX content and SUB data.
    pub fn load_from_data(&mut self, idx_content: &str, sub_data: Vec<u8>) {
        self.dispose();
        self.apply_loaded_data(parse_idx(idx_content), sub_data, true);
    }

    /// Load VobSub from a Matroska subtitle container with embedded S_VOBSUB tracks.
    pub fn load_from_mks(&mut self, mks_data: &[u8]) -> Result<(), String> {
        self.dispose();

        let ExtractedVobSub {
            idx_content,
            sub_data,
            language,
            track_id,
        } = extract_vobsub_from_mks(mks_data)?;

        let mut idx = parse_idx(&idx_content);
        if language.is_some() {
            idx.metadata.language = language;
        }
        if track_id.is_some() {
            idx.metadata.id = track_id;
        }

        self.apply_loaded_data(idx, sub_data, true);
        Ok(())
    }

    /// Load VobSub from SUB file only (scans for timestamps).
    pub fn load_from_sub_only(&mut self, sub_data: Vec<u8>) {
        self.dispose();

        // Default palette
        let palette = VobSubPalette::default();

        // Pre-allocate with estimate (roughly 1 subtitle per 10KB)
        let estimated_count = (sub_data.len() / 10000).max(32);
        let mut timestamps: Vec<VobSubTimestamp> = Vec::with_capacity(estimated_count);
        let mut offset = 0;
        let len = sub_data.len();

        // Look for MPEG-2 PS pack start code
        while offset < len.saturating_sub(4) {
            // Find next potential start code (0x00 0x00 0x01 0xBA)
            if let Some(pos) = memchr(0x00, &sub_data[offset..]) {
                let candidate = offset + pos;

                // Check for full start code: 00 00 01 BA
                if candidate + 3 < len
                    && sub_data[candidate + 1] == 0x00
                    && sub_data[candidate + 2] == 0x01
                    && sub_data[candidate + 3] == 0xBA
                    && let Some((packet, _)) = parse_subtitle_packet(&sub_data, candidate, &palette)
                    && packet.width > 0
                    && packet.height > 0
                {
                    timestamps.push(VobSubTimestamp {
                        timestamp_ms: packet.timestamp_ms,
                        file_position: candidate as u64,
                    });
                }
                offset = candidate + 1;
            } else {
                // No more 0x00 bytes found
                break;
            }
        }

        // Sort and store
        timestamps.sort_by_key(|t| t.timestamp_ms);
        self.timestamps_ms = timestamps.iter().map(|t| t.timestamp_ms).collect();

        let idx = IdxParseResult {
            palette,
            timestamps,
            metadata: Default::default(),
        };
        self.apply_loaded_data(idx, sub_data, false);
    }

    /// Dispose of all resources.
    pub fn dispose(&mut self) {
        self.idx_data = None;
        self.sub_data = None;
        self.timestamps_ms.clear();
        self.packet_cache.clear();
        self.deband_config = DebandConfig::default();
        self.loaded_from_idx = false;
        self.last_render_issue = None;
    }

    /// Get the last non-fatal render issue for diagnostics.
    pub fn last_render_issue(&self) -> String {
        self.last_render_issue.clone().unwrap_or_default()
    }

    /// Get the number of subtitle entries.
    pub fn count(&self) -> usize {
        self.timestamps_ms.len()
    }

    /// Get the presentation width for this subtitle track.
    pub fn screen_width(&self) -> u16 {
        self.idx_data
            .as_ref()
            .map_or(0, |idx_data| idx_data.metadata.width)
    }

    /// Get the presentation height for this subtitle track.
    pub fn screen_height(&self) -> u16 {
        self.idx_data
            .as_ref()
            .map_or(0, |idx_data| idx_data.metadata.height)
    }

    /// Get the declared language code from IDX metadata.
    pub fn language(&self) -> String {
        self.idx_data
            .as_ref()
            .and_then(|idx_data| idx_data.metadata.language.clone())
            .unwrap_or_default()
    }

    /// Get the declared subtitle track ID from IDX metadata.
    pub fn track_id(&self) -> String {
        self.idx_data
            .as_ref()
            .and_then(|idx_data| idx_data.metadata.id.clone())
            .unwrap_or_default()
    }

    /// Check whether IDX metadata was used to load the parser.
    pub fn has_idx_metadata(&self) -> bool {
        self.loaded_from_idx
    }

    /// Get all timestamps in milliseconds.
    pub fn get_timestamps(&self) -> Vec<f64> {
        self.timestamps_ms.iter().map(|&ts| ts as f64).collect()
    }

    /// Find the subtitle index for a given timestamp in milliseconds.
    /// Returns -1 if no subtitle should be displayed at this time.
    pub fn find_index_at_timestamp(&mut self, time_ms: f64) -> i32 {
        if self.timestamps_ms.is_empty() {
            return -1;
        }

        let time_ms_u32 = time_ms as u32;
        let index = binary_search_timestamp(&self.timestamps_ms, time_ms_u32);

        // Get the start time from IDX (what we searched against)
        let start_time = self.timestamps_ms[index];

        // Don't show if we're before this subtitle's start time
        if time_ms_u32 < start_time {
            return -1;
        }

        // Calculate end time
        let end_time = self.calculate_end_time(index, start_time);

        if time_ms_u32 < end_time {
            return index as i32;
        }

        // Current time is past the subtitle's duration
        -1
    }

    /// Get the cue start time in milliseconds.
    pub fn get_cue_start_time(&self, index: usize) -> f64 {
        self.timestamps_ms
            .get(index)
            .copied()
            .map_or(-1.0, |ts| ts as f64)
    }

    /// Get the cue end time in milliseconds.
    pub fn get_cue_end_time(&mut self, index: usize) -> f64 {
        let Some(&start_time) = self.timestamps_ms.get(index) else {
            return -1.0;
        };

        self.calculate_end_time(index, start_time) as f64
    }

    /// Get the cue duration in milliseconds.
    pub fn get_cue_duration(&mut self, index: usize) -> f64 {
        let Some(&start_time) = self.timestamps_ms.get(index) else {
            return -1.0;
        };

        self.calculate_end_time(index, start_time)
            .saturating_sub(start_time) as f64
    }

    /// Get the cue file position in the SUB file.
    pub fn get_cue_file_position(&self, index: usize) -> f64 {
        self.idx_data
            .as_ref()
            .and_then(|idx_data| idx_data.timestamps.get(index).copied())
            .map_or(-1.0, |timestamp| timestamp.file_position as f64)
    }

    fn apply_loaded_data(
        &mut self,
        idx_data: IdxParseResult,
        sub_data: Vec<u8>,
        loaded_from_idx: bool,
    ) {
        self.timestamps_ms = idx_data.timestamps.iter().map(|t| t.timestamp_ms).collect();
        self.idx_data = Some(idx_data);
        self.sub_data = Some(sub_data);
        self.loaded_from_idx = loaded_from_idx;
    }

    /// Calculate the end time for a subtitle at the given index.
    fn calculate_end_time(&mut self, index: usize, start_time: u32) -> u32 {
        // Maximum duration for the last subtitle (no next subtitle to clamp to)
        const MAX_LAST_DURATION_MS: u32 = 5000;

        // Try to get explicit duration from control sequence first
        self.ensure_packet_cached(index);
        let explicit_duration = self
            .cached_packet(index)
            .filter(|p| p.duration_ms > 0 && p.duration_ms != 5000)
            .map(|p| p.duration_ms);

        // Check if we have a next subtitle
        if index + 1 < self.timestamps_ms.len() {
            let next_start = self.timestamps_ms[index + 1];

            if let Some(duration) = explicit_duration {
                let explicit_end = start_time.saturating_add(duration);
                return explicit_end.min(next_start);
            }

            next_start
        } else {
            // Last subtitle - use explicit duration if valid, otherwise default
            if let Some(duration) = explicit_duration {
                return start_time.saturating_add(duration);
            }
            // Default duration for last subtitle
            start_time.saturating_add(MAX_LAST_DURATION_MS)
        }
    }

    fn ensure_packet_cached(&mut self, index: usize) -> Option<()> {
        let idx_data = self.idx_data.as_ref()?;
        if index >= idx_data.timestamps.len() {
            return None;
        }

        if self.packet_cache.contains_key(&index) {
            return Some(());
        }

        let packet = {
            let idx_data = self.idx_data.as_ref()?;
            let sub_data = self.sub_data.as_ref()?;
            let timestamp = idx_data.timestamps.get(index)?;

            parse_subtitle_packet(
                sub_data,
                timestamp.file_position as usize,
                &idx_data.palette,
            )
            .map(|(p, _)| p)
        };

        self.packet_cache.insert(index, packet);
        Some(())
    }

    fn cached_packet(&self, index: usize) -> Option<&SubtitlePacket> {
        self.packet_cache
            .get(&index)
            .and_then(|packet| packet.as_ref())
    }

    /// Render subtitle at the given index and return RGBA data.
    pub fn render_at_index(&mut self, index: usize) -> Option<VobSubFrame> {
        self.last_render_issue = None;

        if index >= self.timestamps_ms.len() {
            self.last_render_issue = Some("INDEX_OUT_OF_RANGE".to_string());
            return None;
        }

        if self.ensure_packet_cached(index).is_none() {
            self.last_render_issue = Some("NO_DATA".to_string());
            return None;
        }

        let Some(idx_data) = self.idx_data.as_ref() else {
            self.last_render_issue = Some("NO_DATA".to_string());
            return None;
        };
        let Some(sub_data) = self.sub_data.as_ref() else {
            self.last_render_issue = Some("NO_DATA".to_string());
            return None;
        };
        let Some(packet) = self.cached_packet(index) else {
            self.last_render_issue = Some("INVALID_PACKET".to_string());
            return None;
        };

        Some(self.render_packet(packet, sub_data, &idx_data.palette, &idx_data.metadata))
    }

    /// Render a packet to a frame.
    fn render_packet(
        &self,
        packet: &SubtitlePacket,
        sub_data: &[u8],
        palette: &VobSubPalette,
        metadata: &super::VobSubMetadata,
    ) -> VobSubFrame {
        let mut rgba = decode_vobsub_rle(packet, sub_data, palette);

        // Apply debanding if enabled
        if self.deband_config.enabled {
            rgba = apply_deband(
                &rgba,
                packet.width as usize,
                packet.height as usize,
                &self.deband_config,
            );
        }

        VobSubFrame {
            screen_width: metadata.width,
            screen_height: metadata.height,
            x: packet.x,
            y: packet.y,
            width: packet.width,
            height: packet.height,
            rgba,
        }
    }

    /// Clear the internal cache.
    pub fn clear_cache(&mut self) {
        self.packet_cache.clear();
    }

    /// Enable or disable debanding.
    pub fn set_deband_enabled(&mut self, enabled: bool) {
        self.deband_config.enabled = enabled;
    }

    /// Set the deband threshold (0.0-255.0, default: 64.0).
    pub fn set_deband_threshold(&mut self, threshold: f32) {
        self.deband_config.threshold = threshold.clamp(0.0, 255.0);
    }

    /// Set the deband sample range in pixels (default: 15).
    pub fn set_deband_range(&mut self, range: u32) {
        self.deband_config.range = range.clamp(1, 64);
    }

    /// Check if debanding is enabled.
    pub fn deband_enabled(&self) -> bool {
        self.deband_config.enabled
    }
}

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

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

    #[test]
    fn dispose_restores_default_deband_config() {
        let mut parser = VobSubParser::new();

        parser.set_deband_enabled(false);
        parser.set_deband_threshold(12.0);
        parser.set_deband_range(3);

        parser.dispose();

        assert!(parser.deband_enabled());
        assert_eq!(
            parser.deband_config.threshold,
            DebandConfig::default().threshold
        );
        assert_eq!(parser.deband_config.range, DebandConfig::default().range);
    }
}

/// A VobSub subtitle frame.
pub struct VobSubFrame {
    pub screen_width: u16,
    pub screen_height: u16,
    pub x: u16,
    pub y: u16,
    pub width: u16,
    pub height: u16,
    pub rgba: Vec<u8>,
}

impl VobSubFrame {
    pub fn screen_width(&self) -> u16 {
        self.screen_width
    }
    pub fn screen_height(&self) -> u16 {
        self.screen_height
    }
    pub fn x(&self) -> u16 {
        self.x
    }
    pub fn y(&self) -> u16 {
        self.y
    }
    pub fn width(&self) -> u16 {
        self.width
    }
    pub fn height(&self) -> u16 {
        self.height
    }

    /// Get RGBA pixel data.
    pub fn get_rgba(&self) -> &[u8] {
        &self.rgba
    }
}