//! Windowed, log_time-ordered k-way merge across the three input MCAPs.
//!
//! Each input is read via its chunk index; only chunks whose `[start,end]`
//! time range overlaps the requested window are considered. Chunks are visited
//! in `message_start_time` order and loaded lazily: a chunk is materialised
//! only once the merge frontier reaches its start time, so peak memory is
//! bounded by the chunk-overlap depth (2 for these recordings), not the file
//! size. Popping the heap minimum yields a globally monotonic stream, filtered
//! to `[window_start, window_end]`.

use std::cmp::Reverse;
use std::collections::BinaryHeap;

use anyhow::Result;
use mcap::records::ChunkIndex;
use mcap::Summary;

/// Which physical camera an input came from — drives topic remapping.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum Role {
    Ego,
    WristLeft,
    WristRight,
}

/// One merged message, detached from the chunk decompression buffer.
pub struct MergedMsg {
    pub role: Role,
    /// Channel id within the source input; topic/schema are looked up in the
    /// catalog built from the input summaries (see pipeline).
    pub in_channel_id: u16,
    pub sequence: u32,
    pub log_time: u64,
    pub publish_time: u64,
    pub data: Vec<u8>,
}

/// A source input plus its role, borrowed for the lifetime of the reader.
pub struct Source<'a> {
    pub role: Role,
    pub mcap: &'a [u8],
    pub summary: &'a Summary,
}

struct PendingChunk<'a> {
    role: Role,
    mcap: &'a [u8],
    summary: &'a Summary,
    ci: ChunkIndex,
}

struct HeapItem {
    log_time: u64,
    /// Global load order — stable tiebreak so equal-`log_time` messages keep
    /// their relative order and heap keys stay unique.
    seq: u64,
    msg: MergedMsg,
}

impl PartialEq for HeapItem {
    fn eq(&self, o: &Self) -> bool {
        (self.log_time, self.seq) == (o.log_time, o.seq)
    }
}
impl Eq for HeapItem {}
impl PartialOrd for HeapItem {
    fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(o))
    }
}
impl Ord for HeapItem {
    fn cmp(&self, o: &Self) -> std::cmp::Ordering {
        (self.log_time, self.seq).cmp(&(o.log_time, o.seq))
    }
}

pub struct WindowReader<'a> {
    start: u64,
    end: u64,
    /// Overlapping chunks across all sources, sorted by `message_start_time`.
    pending: Vec<PendingChunk<'a>>,
    next: usize,
    heap: BinaryHeap<Reverse<HeapItem>>,
    seq: u64,
}

impl<'a> WindowReader<'a> {
    /// `[start, end]` is an inclusive window in absolute `log_time` ns.
    pub fn new(sources: &[Source<'a>], start: u64, end: u64) -> Result<Self> {
        let mut pending = Vec::new();
        for s in sources {
            for ci in &s.summary.chunk_indexes {
                // Keep chunks whose time range intersects the window.
                if ci.message_end_time >= start && ci.message_start_time <= end {
                    pending.push(PendingChunk {
                        role: s.role,
                        mcap: s.mcap,
                        summary: s.summary,
                        ci: ci.clone(),
                    });
                }
            }
        }
        pending.sort_by_key(|p| p.ci.message_start_time);
        Ok(Self {
            start,
            end,
            pending,
            next: 0,
            heap: BinaryHeap::new(),
            seq: 0,
        })
    }

    fn load_chunk(&mut self, idx: usize) -> Result<()> {
        let role = self.pending[idx].role;
        let mcap = self.pending[idx].mcap;
        let summary = self.pending[idx].summary;
        let ci = self.pending[idx].ci.clone();

        // Copy in-window messages out, then drop the decompression buffer
        // before touching the heap.
        let mut owned = Vec::new();
        for m in summary.stream_chunk(mcap, &ci)? {
            let m = m?;
            if m.log_time < self.start || m.log_time > self.end {
                continue;
            }
            owned.push(MergedMsg {
                role,
                in_channel_id: m.channel.id,
                sequence: m.sequence,
                log_time: m.log_time,
                publish_time: m.publish_time,
                data: m.data.into_owned(),
            });
        }
        for msg in owned {
            let seq = self.seq;
            self.seq += 1;
            self.heap.push(Reverse(HeapItem {
                log_time: msg.log_time,
                seq,
                msg,
            }));
        }
        Ok(())
    }
}

impl Iterator for WindowReader<'_> {
    type Item = Result<MergedMsg>;

    fn next(&mut self) -> Option<Self::Item> {
        // Load every chunk that could still contain a message <= the current
        // heap minimum. A chunk whose `message_start_time` exceeds the smallest
        // buffered `log_time` cannot hold an earlier message.
        loop {
            let need_more = match self.heap.peek() {
                None => self.next < self.pending.len(),
                Some(Reverse(top)) => {
                    self.next < self.pending.len()
                        && self.pending[self.next].ci.message_start_time <= top.log_time
                }
            };
            if !need_more {
                break;
            }
            let i = self.next;
            self.next += 1;
            if let Err(e) = self.load_chunk(i) {
                return Some(Err(e));
            }
        }
        self.heap.pop().map(|Reverse(it)| Ok(it.msg))
    }
}
