datum-cdc 0.10.3

use std::collections::HashMap;

use crate::{
    CdcError, CdcResult, ChangeOperation, ColumnMetadata, ColumnValue, RelationMetadata,
    ReplicaIdentity, RowData, TruncateOptions,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct DecodedChange {
    pub(crate) relation: RelationMetadata,
    pub(crate) op: ChangeOperation,
    pub(crate) before: Option<RowData>,
    pub(crate) after: Option<RowData>,
    pub(crate) truncate: Option<TruncateOptions>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct TypeMetadata {
    pub(crate) oid: u32,
    pub(crate) schema: String,
    pub(crate) name: String,
}

#[derive(Debug, Default)]
pub(crate) struct PgOutputDecoder {
    relations: HashMap<u32, RelationMetadata>,
    types: HashMap<u32, TypeMetadata>,
    revision: u64,
}

impl PgOutputDecoder {
    pub(crate) fn decode(&mut self, data: &[u8]) -> CdcResult<Vec<DecodedChange>> {
        let Some((&tag, payload)) = data.split_first() else {
            return Ok(Vec::new());
        };
        let mut reader = PgOutputReader::new(payload);
        match tag {
            b'R' => {
                let relation = reader.read_relation(self.next_revision())?;
                self.relations.insert(relation.oid, relation);
                Ok(Vec::new())
            }
            b'Y' => {
                let ty = reader.read_type()?;
                self.types.insert(ty.oid, ty);
                Ok(Vec::new())
            }
            b'I' => self.decode_insert(&mut reader),
            b'U' => self.decode_update(&mut reader),
            b'D' => self.decode_delete(&mut reader),
            b'T' => self.decode_truncate(&mut reader),
            b'B' | b'C' | b'M' | b'O' => Ok(Vec::new()),
            b'S' | b'E' | b'c' | b'A' | b'K' => Err(CdcError::Parse(format!(
                "unsupported pgoutput streaming/two-phase message tag {:?}",
                tag as char
            ))),
            _ => Err(CdcError::Parse(format!(
                "unsupported pgoutput message tag 0x{tag:02x}"
            ))),
        }
    }

    fn decode_insert(&self, reader: &mut PgOutputReader<'_>) -> CdcResult<Vec<DecodedChange>> {
        let relation = self.read_relation(reader)?;
        reader.expect_tag(b'N', "insert new tuple")?;
        let after = reader.read_tuple(&relation)?;
        reader.finish()?;
        Ok(vec![DecodedChange {
            relation,
            op: ChangeOperation::Insert,
            before: None,
            after: Some(after),
            truncate: None,
        }])
    }

    fn decode_update(&self, reader: &mut PgOutputReader<'_>) -> CdcResult<Vec<DecodedChange>> {
        let relation = self.read_relation(reader)?;
        let tuple_tag = reader.read_u8("update tuple tag")?;
        let before = match tuple_tag {
            b'K' | b'O' => {
                let before = reader.read_tuple(&relation)?;
                reader.expect_tag(b'N', "update new tuple")?;
                Some(before)
            }
            b'N' => None,
            other => {
                return Err(CdcError::Parse(format!(
                    "unexpected update tuple tag 0x{other:02x}"
                )));
            }
        };
        let after = reader.read_tuple(&relation)?;
        reader.finish()?;
        Ok(vec![DecodedChange {
            relation,
            op: ChangeOperation::Update,
            before,
            after: Some(after),
            truncate: None,
        }])
    }

    fn decode_delete(&self, reader: &mut PgOutputReader<'_>) -> CdcResult<Vec<DecodedChange>> {
        let relation = self.read_relation(reader)?;
        let tuple_tag = reader.read_u8("delete tuple tag")?;
        if tuple_tag != b'K' && tuple_tag != b'O' {
            return Err(CdcError::Parse(format!(
                "unexpected delete tuple tag 0x{tuple_tag:02x}"
            )));
        }
        let before = reader.read_tuple(&relation)?;
        reader.finish()?;
        Ok(vec![DecodedChange {
            relation,
            op: ChangeOperation::Delete,
            before: Some(before),
            after: None,
            truncate: None,
        }])
    }

    fn decode_truncate(&self, reader: &mut PgOutputReader<'_>) -> CdcResult<Vec<DecodedChange>> {
        let relation_count = reader.read_i32("truncate relation count")?;
        if relation_count < 0 {
            return Err(CdcError::Parse(format!(
                "negative truncate relation count {relation_count}"
            )));
        }
        let options = reader.read_u8("truncate options")?;
        let truncate = TruncateOptions {
            cascade: options & 0x01 != 0,
            restart_identity: options & 0x02 != 0,
        };
        // Cap the pre-allocation: `relation_count` is a wire-provided i32, and a
        // malformed Truncate message could otherwise reserve gigabytes before the
        // per-relation reads below reject the short body. The vector still grows
        // to the real count when a legitimate message exceeds the cap.
        let mut changes = Vec::with_capacity((relation_count as usize).min(1024));
        for _ in 0..relation_count {
            let relation_id = reader.read_u32("truncate relation id")?;
            let relation = self
                .relations
                .get(&relation_id)
                .cloned()
                .ok_or_else(|| missing_relation(relation_id))?;
            changes.push(DecodedChange {
                relation,
                op: ChangeOperation::Truncate,
                before: None,
                after: None,
                truncate: Some(truncate),
            });
        }
        reader.finish()?;
        Ok(changes)
    }

    fn read_relation(&self, reader: &mut PgOutputReader<'_>) -> CdcResult<RelationMetadata> {
        let relation_id = reader.read_u32("relation id")?;
        self.relations
            .get(&relation_id)
            .cloned()
            .ok_or_else(|| missing_relation(relation_id))
    }

    fn next_revision(&mut self) -> u64 {
        self.revision = self.revision.saturating_add(1);
        self.revision
    }
}

fn missing_relation(relation_id: u32) -> CdcError {
    CdcError::Parse(format!(
        "pgoutput data referenced relation oid {relation_id} before a Relation message"
    ))
}

struct PgOutputReader<'a> {
    bytes: &'a [u8],
    offset: usize,
}

impl<'a> PgOutputReader<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, offset: 0 }
    }

    fn finish(&self) -> CdcResult<()> {
        if self.offset == self.bytes.len() {
            Ok(())
        } else {
            Err(CdcError::Parse(format!(
                "pgoutput message has {} trailing bytes",
                self.bytes.len() - self.offset
            )))
        }
    }

    fn read_relation(&mut self, revision: u64) -> CdcResult<RelationMetadata> {
        let oid = self.read_u32("relation oid")?;
        let schema = self.read_cstring("relation schema")?;
        let table = self.read_cstring("relation table")?;
        let replica_identity = ReplicaIdentity::from_wire(self.read_u8("replica identity")?);
        let column_count = self.read_i16("relation column count")?;
        if column_count < 0 {
            return Err(CdcError::Parse(format!(
                "negative relation column count {column_count}"
            )));
        }
        let mut columns = Vec::with_capacity(column_count as usize);
        for _ in 0..column_count {
            let flags = self.read_u8("relation column flags")?;
            let name = self.read_cstring("relation column name")?;
            let type_oid = self.read_u32("relation column type oid")?;
            let type_modifier = self.read_i32("relation column type modifier")?;
            columns.push(ColumnMetadata {
                name,
                type_oid,
                type_modifier,
                key: flags & 0x01 != 0,
            });
        }
        self.finish()?;
        Ok(RelationMetadata {
            oid,
            schema,
            table,
            replica_identity,
            columns,
            revision,
        })
    }

    fn read_type(&mut self) -> CdcResult<TypeMetadata> {
        let oid = self.read_u32("type oid")?;
        let schema = self.read_cstring("type schema")?;
        let name = self.read_cstring("type name")?;
        self.finish()?;
        Ok(TypeMetadata { oid, schema, name })
    }

    fn read_tuple(&mut self, relation: &RelationMetadata) -> CdcResult<RowData> {
        let column_count = self.read_i16("tuple column count")?;
        if column_count < 0 {
            return Err(CdcError::Parse(format!(
                "negative tuple column count {column_count}"
            )));
        }
        let column_count = usize::try_from(column_count)
            .map_err(|_| CdcError::Parse("tuple column count does not fit usize".into()))?;
        if column_count != relation.columns.len() {
            return Err(CdcError::Parse(format!(
                "tuple for {}.{} has {column_count} columns but relation metadata has {}",
                relation.schema,
                relation.table,
                relation.columns.len()
            )));
        }
        let mut values = Vec::with_capacity(column_count);
        for _ in 0..column_count {
            let value_tag = self.read_u8("tuple value tag")?;
            values.push(match value_tag {
                b'n' => ColumnValue::Null,
                b'u' => ColumnValue::ToastUnchanged,
                b't' => {
                    let value = self.read_len_prefixed("text tuple value")?;
                    ColumnValue::Text(String::from_utf8(value.to_vec()).map_err(|err| {
                        CdcError::Parse(format!("text tuple value is not UTF-8: {err}"))
                    })?)
                }
                b'b' => {
                    let value = self.read_len_prefixed("binary tuple value")?;
                    ColumnValue::Binary(value.to_vec())
                }
                other => {
                    return Err(CdcError::Parse(format!(
                        "unexpected tuple value tag 0x{other:02x}"
                    )));
                }
            });
        }
        Ok(RowData { values })
    }

    fn expect_tag(&mut self, expected: u8, context: &str) -> CdcResult<()> {
        let actual = self.read_u8(context)?;
        if actual == expected {
            Ok(())
        } else {
            Err(CdcError::Parse(format!(
                "{context}: expected tag 0x{expected:02x}, got 0x{actual:02x}"
            )))
        }
    }

    fn read_len_prefixed(&mut self, context: &str) -> CdcResult<&'a [u8]> {
        let len = self.read_i32(context)?;
        if len < 0 {
            return Err(CdcError::Parse(format!("{context}: negative length {len}")));
        }
        let len = usize::try_from(len)
            .map_err(|_| CdcError::Parse(format!("{context}: length does not fit usize")))?;
        self.take(len, context)
    }

    fn read_cstring(&mut self, context: &str) -> CdcResult<String> {
        let remaining = &self.bytes[self.offset..];
        let Some(pos) = remaining.iter().position(|byte| *byte == 0) else {
            return Err(CdcError::Parse(format!(
                "{context}: missing null terminator"
            )));
        };
        let value = &remaining[..pos];
        self.offset += pos + 1;
        String::from_utf8(value.to_vec())
            .map_err(|err| CdcError::Parse(format!("{context}: invalid UTF-8: {err}")))
    }

    fn read_u8(&mut self, context: &str) -> CdcResult<u8> {
        Ok(self.take(1, context)?[0])
    }

    fn read_i16(&mut self, context: &str) -> CdcResult<i16> {
        let bytes = self.take(2, context)?;
        Ok(i16::from_be_bytes([bytes[0], bytes[1]]))
    }

    fn read_i32(&mut self, context: &str) -> CdcResult<i32> {
        let bytes = self.take(4, context)?;
        Ok(i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn read_u32(&mut self, context: &str) -> CdcResult<u32> {
        let bytes = self.take(4, context)?;
        Ok(u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn take(&mut self, len: usize, context: &str) -> CdcResult<&'a [u8]> {
        if self.offset + len > self.bytes.len() {
            return Err(CdcError::Parse(format!(
                "{context}: truncated at byte {} while reading {len} bytes from {} byte message",
                self.offset,
                self.bytes.len()
            )));
        }
        let start = self.offset;
        self.offset += len;
        Ok(&self.bytes[start..self.offset])
    }
}

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

    const RELATION: &[u8] = &[
        b'R', 0, 0, 0, 42, b'p', b'u', b'b', b'l', b'i', b'c', 0, b'c', b'd', b'c', b'_', b'e',
        b'v', b'e', b'n', b't', b's', 0, b'f', 0, 3, 1, b'i', b'd', 0, 0, 0, 0, 20, 0xff, 0xff,
        0xff, 0xff, 0, b'r', b'u', b'n', b'_', b'i', b'd', 0, 0, 0, 0, 25, 0xff, 0xff, 0xff, 0xff,
        0, b'v', b'a', b'l', b'u', b'e', 0, 0, 0, 0, 20, 0xff, 0xff, 0xff, 0xff,
    ];
    const INSERT: &[u8] = &[
        b'I', 0, 0, 0, 42, b'N', 0, 3, b't', 0, 0, 0, 1, b'7', b't', 0, 0, 0, 5, b'r', b'u', b'n',
        b'-', b'1', b't', 0, 0, 0, 2, b'4', b'2',
    ];
    const UPDATE: &[u8] = &[
        b'U', 0, 0, 0, 42, b'O', 0, 3, b't', 0, 0, 0, 1, b'7', b't', 0, 0, 0, 5, b'r', b'u', b'n',
        b'-', b'1', b't', 0, 0, 0, 2, b'4', b'2', b'N', 0, 3, b't', 0, 0, 0, 1, b'7', b't', 0, 0,
        0, 5, b'r', b'u', b'n', b'-', b'1', b't', 0, 0, 0, 2, b'4', b'3',
    ];
    const DELETE: &[u8] = &[
        b'D', 0, 0, 0, 42, b'O', 0, 3, b't', 0, 0, 0, 1, b'7', b't', 0, 0, 0, 5, b'r', b'u', b'n',
        b'-', b'1', b't', 0, 0, 0, 2, b'4', b'3',
    ];
    const TRUNCATE: &[u8] = &[b'T', 0, 0, 0, 1, 0x03, 0, 0, 0, 42];

    #[test]
    fn decodes_recorded_relation_insert_update_delete_truncate_fixtures() {
        let mut decoder = PgOutputDecoder::default();
        assert!(decoder.decode(RELATION).unwrap().is_empty());

        let insert = decoder.decode(INSERT).unwrap();
        assert_eq!(insert.len(), 1);
        assert_eq!(insert[0].op, ChangeOperation::Insert);
        assert_eq!(
            insert[0]
                .after
                .as_ref()
                .unwrap()
                .get_text(&insert[0].relation, "value"),
            Some("42")
        );

        let update = decoder.decode(UPDATE).unwrap();
        assert_eq!(update[0].op, ChangeOperation::Update);
        assert_eq!(
            update[0]
                .before
                .as_ref()
                .unwrap()
                .get_text(&update[0].relation, "value"),
            Some("42")
        );
        assert_eq!(
            update[0]
                .after
                .as_ref()
                .unwrap()
                .get_text(&update[0].relation, "value"),
            Some("43")
        );

        let delete = decoder.decode(DELETE).unwrap();
        assert_eq!(delete[0].op, ChangeOperation::Delete);
        assert_eq!(
            delete[0]
                .before
                .as_ref()
                .unwrap()
                .get_text(&delete[0].relation, "value"),
            Some("43")
        );

        let truncate = decoder.decode(TRUNCATE).unwrap();
        assert_eq!(truncate[0].op, ChangeOperation::Truncate);
        assert_eq!(
            truncate[0].truncate,
            Some(TruncateOptions {
                cascade: true,
                restart_identity: true,
            })
        );
    }

    #[test]
    fn decodes_null_and_toast_unchanged_values() {
        let mut decoder = PgOutputDecoder::default();
        assert!(decoder.decode(RELATION).unwrap().is_empty());

        // INSERT with a NULL ('n') in the id and run_id columns and text "42" in value.
        const INSERT_WITH_NULL: &[u8] = &[
            b'I', 0, 0, 0, 42, b'N', 0, 3, b'n', b'n', b't', 0, 0, 0, 2, b'4', b'2',
        ];
        let insert = decoder.decode(INSERT_WITH_NULL).unwrap();
        assert_eq!(insert.len(), 1);
        let after = insert[0].after.as_ref().unwrap();
        assert_eq!(
            after.get(&insert[0].relation, "id"),
            Some(&ColumnValue::Null)
        );
        assert_eq!(after.get_text(&insert[0].relation, "value"), Some("42"));

        // UPDATE (key unchanged, so 'N'-only) whose run_id column is an unchanged
        // TOAST value ('u') that pgoutput does not resend.
        const UPDATE_WITH_TOAST: &[u8] = &[
            b'U', 0, 0, 0, 42, b'N', 0, 3, b't', 0, 0, 0, 1, b'7', b'u', b't', 0, 0, 0, 2, b'4',
            b'3',
        ];
        let update = decoder.decode(UPDATE_WITH_TOAST).unwrap();
        assert_eq!(update.len(), 1);
        assert!(update[0].before.is_none());
        let after = update[0].after.as_ref().unwrap();
        assert_eq!(
            after.get(&update[0].relation, "run_id"),
            Some(&ColumnValue::ToastUnchanged)
        );
        // A TOAST-unchanged column has no readable text.
        assert_eq!(after.get_text(&update[0].relation, "run_id"), None);
        assert_eq!(after.get_text(&update[0].relation, "value"), Some("43"));
    }

    #[test]
    fn relation_must_arrive_before_data() {
        let mut decoder = PgOutputDecoder::default();
        assert!(matches!(
            decoder.decode(INSERT),
            Err(CdcError::Parse(message)) if message.contains("before a Relation message")
        ));
    }

    #[test]
    fn truncate_with_huge_claimed_count_errors_without_allocating() {
        let mut decoder = PgOutputDecoder::default();
        // Truncate claiming i32::MAX relations with an empty body. The claimed
        // count must not be trusted for pre-allocation (multi-GB reservation);
        // decoding must fail on the missing relation id instead.
        const TRUNCATE_HUGE: &[u8] = &[b'T', 0x7f, 0xff, 0xff, 0xff, 0];
        assert!(matches!(
            decoder.decode(TRUNCATE_HUGE),
            Err(CdcError::Parse(message)) if message.contains("truncate relation id")
        ));
    }

    #[test]
    fn tolerates_origin_message_and_keeps_decoding() {
        let mut decoder = PgOutputDecoder::default();
        // An Origin ('O') message appears when the WAL carries changes with a
        // replication origin set (cascaded/bi-directional logical replication).
        // Like Begin/Commit/Message it carries no row data and must be skipped,
        // not rejected — otherwise the source stalls in a reconnect loop.
        // Body: origin commit LSN (int64) + origin name (cstring).
        const ORIGIN: &[u8] = &[b'O', 0, 0, 0, 0, 0, 0, 0, 0x10, b'p', b'g', b'_', b'1', 0];
        assert!(decoder.decode(ORIGIN).unwrap().is_empty());

        // The decoder still handles a following relation + insert normally.
        assert!(decoder.decode(RELATION).unwrap().is_empty());
        let insert = decoder.decode(INSERT).unwrap();
        assert_eq!(insert.len(), 1);
        assert_eq!(insert[0].op, ChangeOperation::Insert);
    }
}