use std::path::{Path, PathBuf};
use std::sync::Arc;

use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncSeekExt, BufReader};

use serde::{Deserialize, Serialize};

use crate::error::Result;
use crate::types::{
    MessageEnvelope, MessageId, NamespaceId, PublishedMessage, TopicName, TopicOffset,
};

/// Legacy envelope format (v1) — used for deserializing old segment records
/// that were written before the `headers` field was added.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct MessageEnvelopeV1 {
    pub id: MessageId,
    pub namespace: NamespaceId,
    pub topic: TopicName,
    pub payload: Vec<u8>,
    pub timestamp: u64,
}

impl MessageEnvelopeV1 {
    /// Convert to the current `MessageEnvelope` with empty headers and no TTL.
    fn into_current(self) -> MessageEnvelope {
        MessageEnvelope::new(
            self.id,
            self.namespace,
            self.topic,
            self.payload,
            self.timestamp,
        )
    }
}

/// Streaming iterator over records in a segment file.
///
/// Reads one record at a time via [`next`](Self::next). Stops at EOF or first
/// corruption (same recovery semantics as `replay_segment`).
pub struct RecordStream {
    reader: BufReader<fs::File>,
    path: PathBuf,
    /// Byte position of the record that will be returned by the next `next()` call.
    byte_pos: u64,
}

impl RecordStream {
    /// Open a segment file for streaming record reads.
    pub async fn open(path: &Path) -> Result<Self> {
        let file = fs::File::open(path).await?;
        Ok(Self {
            reader: BufReader::new(file),
            path: path.to_path_buf(),
            byte_pos: 0,
        })
    }

    /// Current byte position within the segment file.
    pub fn position(&self) -> u64 {
        self.byte_pos
    }

    /// Read the next record. Returns `Ok(None)` at EOF or on corruption.
    pub async fn next(&mut self) -> Result<Option<(TopicOffset, MessageEnvelope)>> {
        // Read 4-byte record length
        let mut len_buf = [0u8; 4];
        match self.reader.read_exact(&mut len_buf).await {
            Ok(_) => {}
            // Clean EOF at a record boundary — normal end of segment
            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(e.into()),
        }

        let record_len = u32::from_be_bytes(len_buf) as usize;

        if record_len < 8 {
            tracing::warn!(
                "invalid record_len={} in segment {:?}",
                record_len,
                self.path
            );
            return Ok(None);
        }

        // Read the record body
        let mut record_buf = vec![0u8; record_len];
        match self.reader.read_exact(&mut record_buf).await {
            Ok(_) => {}
            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
                tracing::warn!(
                    "truncated record body (record_len={}) in segment {:?}",
                    record_len,
                    self.path
                );
                return Ok(None);
            }
            Err(e) => return Err(e.into()),
        }

        let record_start = self.byte_pos;
        self.byte_pos = record_start + 4 + record_len as u64;

        let offset = TopicOffset(u64::from_be_bytes(record_buf[..8].try_into().unwrap()));

        // Detect record format version.
        // v2 records start with [0x00][version] after the 8-byte offset.
        // v1 (legacy) records start with bincode(MessageEnvelopeV1) where the
        // first byte is the LE low byte of MessageId (always ≥ 1, so non-zero).
        let envelope_data = &record_buf[8..];
        let result = if !envelope_data.is_empty() && envelope_data[0] == 0x00 {
            // v2+ format: skip sentinel (0x00) and version byte
            if envelope_data.len() < 2 {
                tracing::warn!("truncated v2 header in segment {:?}", self.path);
                return Ok(None);
            }
            // envelope_data[1] is the version; currently only 0x02
            bincode::deserialize::<MessageEnvelope>(&envelope_data[2..])
        } else {
            // v1 (legacy) format: deserialize as old envelope, convert
            match bincode::deserialize::<MessageEnvelopeV1>(envelope_data) {
                Ok(v1) => Ok(v1.into_current()),
                Err(e) => Err(e),
            }
        };

        match result {
            Ok(envelope) => Ok(Some((offset, envelope))),
            Err(e) => {
                let msg = e.to_string();
                if msg.contains("SystemTime") || msg.contains("invalid type") {
                    tracing::error!(
                        "segment {:?} uses an incompatible format \
                         (pre-u64-timestamp). Delete the data directory and restart.",
                        self.path
                    );
                } else {
                    tracing::warn!("corrupt envelope in segment {:?}: {}", self.path, e);
                }
                Ok(None)
            }
        }
    }

    /// Byte position of the record that was just returned by `next()`.
    /// Call this right after `next()` to get the start position of that record.
    /// Returns the position *before* the record (i.e., `position() - record_size`).
    pub fn last_record_position(&self, record_total_bytes: u64) -> u64 {
        self.byte_pos - record_total_bytes
    }

    /// Seek the underlying reader to a byte position within the segment.
    /// Subsequent `next()` calls read from this position.
    pub async fn seek_to_position(&mut self, pos: u64) -> Result<()> {
        self.reader.seek(std::io::SeekFrom::Start(pos)).await?;
        self.byte_pos = pos;
        Ok(())
    }
}

/// Replay all messages from segment files in the given directory, in offset order.
pub async fn replay_topic(seg_dir: &Path) -> Result<Vec<PublishedMessage>> {
    let segments = list_segments(seg_dir).await?;
    let mut messages = Vec::new();
    for seg_path in segments {
        let mut seg_msgs = replay_segment(&seg_path).await?;
        messages.append(&mut seg_msgs);
    }
    messages.sort_by_key(|m| m.offset);
    Ok(messages)
}

/// List all `*.seg` files in the directory, sorted by filename (which encodes the base offset).
/// Returns an empty Vec if the directory doesn't exist.
pub async fn list_segments(seg_dir: &Path) -> Result<Vec<PathBuf>> {
    if !seg_dir.exists() {
        return Ok(Vec::new());
    }

    let mut entries = Vec::new();
    let mut read_dir = fs::read_dir(seg_dir).await?;
    while let Some(entry) = read_dir.next_entry().await? {
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) == Some("seg") {
            entries.push(path);
        }
    }
    entries.sort();
    Ok(entries)
}

/// Replay a single segment file using streaming reads via [`RecordStream`].
/// Reads records incrementally instead of loading the entire file into memory.
/// Stops at first corruption or truncation.
///
/// **Format note:** Segments written before the `u64` timestamp migration
/// (when `MessageEnvelope.timestamp` was `SystemTime`) are not compatible.
/// If replay fails with deserialization errors, delete the data directory
/// and restart.
async fn replay_segment(path: &Path) -> Result<Vec<PublishedMessage>> {
    let mut stream = RecordStream::open(path).await?;
    let mut messages = Vec::new();
    while let Some((offset, envelope)) = stream.next().await? {
        messages.push(PublishedMessage {
            envelope: Arc::new(envelope),
            offset,
        });
    }
    Ok(messages)
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::types::{MessageId, NamespaceId, TopicName, now_millis};
    use tempfile::TempDir;
    use tokio::fs;

    pub fn make_msg(offset: u64) -> PublishedMessage {
        PublishedMessage {
            envelope: Arc::new(MessageEnvelope::new(
                MessageId(offset + 1),
                NamespaceId(format!("ns-{offset}")),
                TopicName(format!("topic-{offset}")),
                format!("payload-{offset}").into_bytes(),
                now_millis(),
            )),
            offset: TopicOffset(offset),
        }
    }

    pub fn encode_record(msg: &PublishedMessage) -> Vec<u8> {
        let envelope_bytes = bincode::serialize(&msg.envelope).unwrap();
        let record_len = (8 + 2 + envelope_bytes.len()) as u32; // +2 for sentinel + version
        let mut buf = Vec::new();
        buf.extend_from_slice(&record_len.to_be_bytes());
        buf.extend_from_slice(&msg.offset.0.to_be_bytes());
        buf.push(0x00); // sentinel
        buf.push(0x02); // version
        buf.extend_from_slice(&envelope_bytes);
        buf
    }

    /// Encode a record in legacy v1 format (no version prefix, no headers).
    /// Used for testing backward compatibility with old segments.
    pub fn encode_record_v1(msg: &PublishedMessage) -> Vec<u8> {
        let v1 = MessageEnvelopeV1 {
            id: msg.envelope.id,
            namespace: msg.envelope.namespace.clone(),
            topic: msg.envelope.topic.clone(),
            payload: msg.envelope.payload.clone(),
            timestamp: msg.envelope.timestamp,
        };
        let envelope_bytes = bincode::serialize(&v1).unwrap();
        let record_len = (8 + envelope_bytes.len()) as u32;
        let mut buf = Vec::new();
        buf.extend_from_slice(&record_len.to_be_bytes());
        buf.extend_from_slice(&msg.offset.0.to_be_bytes());
        buf.extend_from_slice(&envelope_bytes);
        buf
    }

    #[tokio::test]
    async fn replay_empty_dir() {
        let tmp = TempDir::new().unwrap();
        let result = replay_topic(tmp.path()).await.unwrap();
        assert!(result.is_empty());
    }

    #[tokio::test]
    async fn replay_nonexistent_dir() {
        let tmp = TempDir::new().unwrap();
        let nonexistent = tmp.path().join("does-not-exist");
        let result = replay_topic(&nonexistent).await.unwrap();
        assert!(result.is_empty());
    }

    #[tokio::test]
    async fn replay_single_segment() {
        let tmp = TempDir::new().unwrap();
        let mut data = Vec::new();
        for i in 0..5 {
            data.extend_from_slice(&encode_record(&make_msg(i)));
        }
        fs::write(tmp.path().join("00000000000000000000.seg"), &data)
            .await
            .unwrap();

        let result = replay_topic(tmp.path()).await.unwrap();
        assert_eq!(result.len(), 5);
        for (i, msg) in result.iter().enumerate() {
            assert_eq!(msg.offset, TopicOffset(i as u64));
            assert_eq!(msg.envelope.payload, format!("payload-{i}").into_bytes());
        }
    }

    #[tokio::test]
    async fn replay_multiple_segments_in_order() {
        let tmp = TempDir::new().unwrap();

        // Segment 1: offsets 0-2
        let mut seg1 = Vec::new();
        for i in 0..3 {
            seg1.extend_from_slice(&encode_record(&make_msg(i)));
        }
        fs::write(tmp.path().join("00000000000000000000.seg"), &seg1)
            .await
            .unwrap();

        // Segment 2: offsets 3-4
        let mut seg2 = Vec::new();
        for i in 3..5 {
            seg2.extend_from_slice(&encode_record(&make_msg(i)));
        }
        fs::write(tmp.path().join("00000000000000000003.seg"), &seg2)
            .await
            .unwrap();

        let result = replay_topic(tmp.path()).await.unwrap();
        assert_eq!(result.len(), 5);
        for (i, msg) in result.iter().enumerate() {
            assert_eq!(msg.offset, TopicOffset(i as u64));
        }
    }

    #[tokio::test]
    async fn replay_truncated_record_recovers() {
        let tmp = TempDir::new().unwrap();
        let mut data = Vec::new();
        for i in 0..5 {
            data.extend_from_slice(&encode_record(&make_msg(i)));
        }
        // Truncate the last record by removing some bytes
        data.truncate(data.len() - 20);

        fs::write(tmp.path().join("00000000000000000000.seg"), &data)
            .await
            .unwrap();

        let result = replay_topic(tmp.path()).await.unwrap();
        // Should recover at least the first 4 records
        assert!(result.len() >= 4);
        assert!(result.len() < 5);
        for (i, msg) in result.iter().enumerate() {
            assert_eq!(msg.offset, TopicOffset(i as u64));
        }
    }

    #[tokio::test]
    async fn replay_ignores_non_seg_files() {
        let tmp = TempDir::new().unwrap();

        // Write a valid segment
        let mut data = Vec::new();
        data.extend_from_slice(&encode_record(&make_msg(0)));
        fs::write(tmp.path().join("00000000000000000000.seg"), &data)
            .await
            .unwrap();

        // Write some non-seg files that should be ignored
        fs::write(tmp.path().join("notes.txt"), b"hello")
            .await
            .unwrap();
        fs::write(tmp.path().join("data.log"), b"log data")
            .await
            .unwrap();

        let result = replay_topic(tmp.path()).await.unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].offset, TopicOffset(0));
    }

    #[tokio::test]
    async fn record_stream_reads_all_records() {
        let tmp = TempDir::new().unwrap();
        let mut data = Vec::new();
        for i in 0..5 {
            data.extend_from_slice(&encode_record(&make_msg(i)));
        }
        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let mut records = Vec::new();
        while let Some((offset, envelope)) = stream.next().await.unwrap() {
            records.push((offset, envelope));
        }

        assert_eq!(records.len(), 5);
        for (i, (offset, envelope)) in records.iter().enumerate() {
            assert_eq!(*offset, TopicOffset(i as u64));
            assert_eq!(envelope.payload, format!("payload-{i}").into_bytes());
        }
    }

    #[tokio::test]
    async fn record_stream_stops_on_corruption() {
        let tmp = TempDir::new().unwrap();
        let mut data = Vec::new();
        for i in 0..5 {
            data.extend_from_slice(&encode_record(&make_msg(i)));
        }
        // Truncate the last record
        data.truncate(data.len() - 20);

        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let mut count = 0;
        while stream.next().await.unwrap().is_some() {
            count += 1;
        }
        // Should read at least the first 4 records, stop at truncated 5th
        assert!(count >= 4);
        assert!(count < 5);
    }

    #[tokio::test]
    async fn v2_record_roundtrip_with_headers() {
        use crate::types::Header;
        let tmp = TempDir::new().unwrap();

        let mut msg = make_msg(0);
        Arc::make_mut(&mut msg.envelope).headers = vec![
            Header {
                key: b"Content-Type".to_vec(),
                value: b"application/json".to_vec(),
            },
            Header {
                key: b"Correlation-Id".to_vec(),
                value: b"req-12345".to_vec(),
            },
        ];

        let mut data = Vec::new();
        data.extend_from_slice(&encode_record(&msg));

        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let (offset, envelope) = stream.next().await.unwrap().unwrap();

        assert_eq!(offset, TopicOffset(0));
        assert_eq!(envelope.headers.len(), 2);
        assert_eq!(envelope.headers[0].key, b"Content-Type");
        assert_eq!(envelope.headers[0].value, b"application/json");
        assert_eq!(envelope.headers[1].key, b"Correlation-Id");
        assert_eq!(envelope.headers[1].value, b"req-12345");
        assert_eq!(envelope.payload, msg.envelope.payload);
    }

    #[tokio::test]
    async fn v1_record_reads_with_empty_headers() {
        let tmp = TempDir::new().unwrap();

        // Write records in legacy v1 format (no version prefix, no headers)
        let mut data = Vec::new();
        for i in 0..3 {
            data.extend_from_slice(&encode_record_v1(&make_msg(i)));
        }

        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let mut records = Vec::new();
        while let Some((offset, envelope)) = stream.next().await.unwrap() {
            records.push((offset, envelope));
        }

        assert_eq!(records.len(), 3);
        for (i, (offset, envelope)) in records.iter().enumerate() {
            assert_eq!(*offset, TopicOffset(i as u64));
            assert_eq!(envelope.payload, format!("payload-{i}").into_bytes());
            // v1 records should have empty headers
            assert!(
                envelope.headers.is_empty(),
                "v1 records should have no headers"
            );
        }
    }

    #[tokio::test]
    async fn mixed_v1_v2_segment_readable() {
        use crate::types::Header;
        let tmp = TempDir::new().unwrap();

        let mut data = Vec::new();

        // First two records in v1 format (no headers)
        data.extend_from_slice(&encode_record_v1(&make_msg(0)));
        data.extend_from_slice(&encode_record_v1(&make_msg(1)));

        // Next two records in v2 format (with headers)
        let mut msg2 = make_msg(2);
        Arc::make_mut(&mut msg2.envelope).headers = vec![Header {
            key: b"trace-id".to_vec(),
            value: b"abc123".to_vec(),
        }];
        data.extend_from_slice(&encode_record(&msg2));
        data.extend_from_slice(&encode_record(&make_msg(3))); // v2, empty headers

        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let mut records = Vec::new();
        while let Some((offset, envelope)) = stream.next().await.unwrap() {
            records.push((offset, envelope));
        }

        assert_eq!(records.len(), 4);

        // v1 records: no headers
        assert!(records[0].1.headers.is_empty());
        assert!(records[1].1.headers.is_empty());

        // v2 record with header
        assert_eq!(records[2].1.headers.len(), 1);
        assert_eq!(records[2].1.headers[0].key, b"trace-id");
        assert_eq!(records[2].1.headers[0].value, b"abc123");

        // v2 record with empty headers
        assert!(records[3].1.headers.is_empty());

        // All offsets sequential
        for (i, (offset, _)) in records.iter().enumerate() {
            assert_eq!(*offset, TopicOffset(i as u64));
        }
    }

    #[tokio::test]
    async fn v2_record_empty_headers_roundtrip() {
        let tmp = TempDir::new().unwrap();

        // v2 format with zero headers should still work
        let msg = make_msg(0);
        assert!(msg.envelope.headers.is_empty());

        let mut data = Vec::new();
        data.extend_from_slice(&encode_record(&msg));

        let seg_path = tmp.path().join("00000000000000000000.seg");
        fs::write(&seg_path, &data).await.unwrap();

        let mut stream = RecordStream::open(&seg_path).await.unwrap();
        let (offset, envelope) = stream.next().await.unwrap().unwrap();

        assert_eq!(offset, TopicOffset(0));
        assert!(envelope.headers.is_empty());
        assert_eq!(envelope.payload, msg.envelope.payload);
    }
}
