pgdumpx 0.2.0

use crate::{
    Limits,
    error::PgDumpError,
    metadata_budget::MetadataBudget,
    model::{ArchiveString, DumpId},
};
use std::{
    collections::{HashMap, HashSet},
    str::Utf8Error,
};

/// The logical representation recorded for one table-data entry.
///
/// [`TableDataRepresentation::CopyText`] is the v0.1 row-aware path.
/// [`TableDataRepresentation::Insert`], [`TableDataRepresentation::Binary`], and
/// [`TableDataRepresentation::Other`] can still be meaningful for metadata or raw
/// selected-entry access, but [`crate::Archive::table_rows`] rejects them with
/// [`PgDumpError::UnsupportedTableDataRepresentation`] rather than guessing a COPY
/// parser for incompatible bytes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum TableDataRepresentation {
    /// PostgreSQL COPY text data with byte-oriented column metadata.
    CopyText,
    /// SQL INSERT statements generated by `pg_dump --inserts` variants.
    Insert,
    /// PostgreSQL binary COPY data, which v0.1 does not decode as text rows.
    Binary,
    /// Another row-aware representation not supported by the v0.1 COPY parser.
    Other,
}

/// One ordered PostgreSQL COPY column name derived from TOC metadata.
///
/// Column names are stored as logical identifier bytes from the recorded
/// `COPY ... FROM stdin` statement and do not require UTF-8. Their position in the
/// slice returned by [`crate::TableRef::columns`] or [`crate::TableRowReader::columns`]
/// is the corresponding field index in parsed rows.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Column {
    name: ArchiveString,
}

impl Column {
    fn new(name: Vec<u8>) -> Self {
        Self {
            name: ArchiveString::from_bytes(name),
        }
    }

    /// Returns the exact logical identifier bytes from the recorded COPY statement.
    pub fn name_bytes(&self) -> &[u8] {
        self.name.as_bytes()
    }

    /// Interprets the column name as UTF-8 without changing the stored bytes.
    ///
    /// Invalid UTF-8 is reported by [`Utf8Error`]; it does not make the byte-oriented
    /// column metadata itself invalid.
    pub fn name_str(&self) -> Result<&str, Utf8Error> {
        self.name.to_str()
    }
}

#[derive(Debug)]
pub(crate) enum TableDataMetadata {
    Copy(CopyColumnLayout),
    Unavailable,
    Malformed { reason: &'static str },
    Unsupported(TableDataRepresentation),
}

impl TableDataMetadata {
    pub(crate) fn validate_row_access(&self, dump_id: DumpId) -> Result<(), PgDumpError> {
        match self {
            Self::Unsupported(representation) => {
                Err(PgDumpError::UnsupportedTableDataRepresentation {
                    dump_id: dump_id.as_i32(),
                    representation: *representation,
                })
            }
            Self::Copy(_) | Self::Unavailable | Self::Malformed { .. } => Ok(()),
        }
    }

    pub(crate) fn representation(
        &self,
        dump_id: DumpId,
    ) -> Result<TableDataRepresentation, PgDumpError> {
        match self {
            Self::Copy(_) => Ok(TableDataRepresentation::CopyText),
            Self::Unsupported(representation) => Ok(*representation),
            Self::Unavailable => Err(PgDumpError::CopyColumnMetadataUnavailable {
                dump_id: dump_id.as_i32(),
            }),
            Self::Malformed { reason } => Err(PgDumpError::MalformedCopyStatement {
                dump_id: dump_id.as_i32(),
                reason,
            }),
        }
    }

    pub(crate) fn columns(&self, dump_id: DumpId) -> Result<&[Column], PgDumpError> {
        match self {
            Self::Copy(layout) => Ok(&layout.columns),
            Self::Unavailable => Err(PgDumpError::CopyColumnMetadataUnavailable {
                dump_id: dump_id.as_i32(),
            }),
            Self::Malformed { reason } => Err(PgDumpError::MalformedCopyStatement {
                dump_id: dump_id.as_i32(),
                reason,
            }),
            Self::Unsupported(representation) => {
                Err(PgDumpError::UnsupportedTableDataRepresentation {
                    dump_id: dump_id.as_i32(),
                    representation: *representation,
                })
            }
        }
    }

    pub(crate) fn column_index(
        &self,
        dump_id: DumpId,
        name: &[u8],
    ) -> Result<Option<usize>, PgDumpError> {
        match self {
            Self::Copy(layout) => Ok(layout.by_name.get(name).copied()),
            Self::Unavailable => Err(PgDumpError::CopyColumnMetadataUnavailable {
                dump_id: dump_id.as_i32(),
            }),
            Self::Malformed { reason } => Err(PgDumpError::MalformedCopyStatement {
                dump_id: dump_id.as_i32(),
                reason,
            }),
            Self::Unsupported(representation) => {
                Err(PgDumpError::UnsupportedTableDataRepresentation {
                    dump_id: dump_id.as_i32(),
                    representation: *representation,
                })
            }
        }
    }
}

#[derive(Debug)]
pub(crate) struct CopyColumnLayout {
    columns: Vec<Column>,
    by_name: HashMap<Vec<u8>, usize>,
}

#[cfg(test)]
#[allow(dead_code)]
pub(crate) fn parse_table_data_metadata(
    dump_id: DumpId,
    copy_statement: Option<&[u8]>,
) -> Result<TableDataMetadata, PgDumpError> {
    parse_table_data_metadata_with_limits(dump_id, copy_statement, Limits::default())
}

#[cfg(test)]
pub(crate) fn parse_table_data_metadata_with_limits(
    dump_id: DumpId,
    copy_statement: Option<&[u8]>,
    limits: Limits,
) -> Result<TableDataMetadata, PgDumpError> {
    let mut budget = MetadataBudget::new(limits)?;
    parse_table_data_metadata_with_limits_and_budget(dump_id, copy_statement, limits, &mut budget)
}

pub(crate) fn parse_table_data_metadata_with_limits_and_budget(
    dump_id: DumpId,
    copy_statement: Option<&[u8]>,
    limits: Limits,
    budget: &mut MetadataBudget,
) -> Result<TableDataMetadata, PgDumpError> {
    let Some(statement) = copy_statement else {
        return Ok(TableDataMetadata::Unsupported(
            TableDataRepresentation::Insert,
        ));
    };
    if statement.is_empty() {
        return Ok(TableDataMetadata::Unavailable);
    }

    let parsed =
        match CopyStatementParser::new(statement, dump_id, limits.max_fields_per_row()).parse() {
            Ok(parsed) => parsed,
            Err(MetadataParseError::Malformed(reason)) => {
                return Ok(TableDataMetadata::Malformed { reason });
            }
            Err(MetadataParseError::Fatal(error)) => return Err(error),
        };

    match parsed {
        ParsedStatement::Copy(columns) => match CopyColumnLayout::new(columns, dump_id, budget)? {
            Ok(layout) => Ok(TableDataMetadata::Copy(layout)),
            Err(reason) => Ok(TableDataMetadata::Malformed { reason }),
        },
        ParsedStatement::Unsupported(representation) => {
            Ok(TableDataMetadata::Unsupported(representation))
        }
    }
}

impl CopyColumnLayout {
    fn new(
        columns: Vec<Column>,
        dump_id: DumpId,
        budget: &mut MetadataBudget,
    ) -> Result<Result<Self, &'static str>, PgDumpError> {
        let requested = to_u64(columns.len())?;
        let mut unique_names = HashSet::new();
        unique_names.try_reserve(columns.len()).map_err(|_| {
            PgDumpError::CopyColumnMetadataAllocationFailed {
                dump_id: dump_id.as_i32(),
                requested,
            }
        })?;

        let mut logical_name_bytes = 0_usize;
        for column in &columns {
            if !unique_names.insert(column.name_bytes()) {
                return Ok(Err("duplicate COPY column name"));
            }
            logical_name_bytes = logical_name_bytes
                .checked_add(column.name_bytes().len())
                .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
        }

        let retained_name_bytes = logical_name_bytes
            .checked_mul(2)
            .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
        budget.charge_index_bytes(retained_name_bytes, "COPY column metadata and lookup")?;

        let mut by_name = HashMap::new();
        by_name.try_reserve(columns.len()).map_err(|_| {
            PgDumpError::CopyColumnMetadataAllocationFailed {
                dump_id: dump_id.as_i32(),
                requested,
            }
        })?;

        for (index, column) in columns.iter().enumerate() {
            let key = clone_bytes(column.name_bytes(), dump_id)?;
            by_name.insert(key, index);
        }

        Ok(Ok(Self { columns, by_name }))
    }
}

enum ParsedStatement {
    Copy(Vec<Column>),
    Unsupported(TableDataRepresentation),
}

enum MetadataParseError {
    Malformed(&'static str),
    Fatal(PgDumpError),
}

impl From<PgDumpError> for MetadataParseError {
    fn from(error: PgDumpError) -> Self {
        Self::Fatal(error)
    }
}

struct CopyStatementParser<'a> {
    input: &'a [u8],
    position: usize,
    dump_id: DumpId,
    max_columns: usize,
}

impl<'a> CopyStatementParser<'a> {
    const fn new(input: &'a [u8], dump_id: DumpId, max_columns: usize) -> Self {
        Self {
            input,
            position: 0,
            dump_id,
            max_columns,
        }
    }

    fn parse(mut self) -> Result<ParsedStatement, MetadataParseError> {
        self.expect_bytes(b"COPY", "COPY statement must begin with COPY")?;
        self.require_whitespace("COPY must be followed by a relation name")?;
        self.skip_relation()?;

        let separated = self.skip_whitespace();
        let columns = if self.consume_byte(b'(') {
            self.parse_columns()?
        } else {
            if !separated {
                return Err(MetadataParseError::Malformed(
                    "relation name must be separated from FROM",
                ));
            }
            Vec::new()
        };

        self.skip_whitespace();
        self.expect_bytes(b"FROM", "COPY column list must be followed by FROM")?;
        self.require_whitespace("FROM must be followed by stdin")?;
        self.expect_bytes(b"stdin", "pg_dump COPY source must be stdin")?;
        self.skip_whitespace();

        if self.consume_byte(b';') {
            self.skip_whitespace();
            if self.position == self.input.len() {
                return Ok(ParsedStatement::Copy(columns));
            }
            return Err(MetadataParseError::Malformed(
                "unexpected bytes after COPY statement terminator",
            ));
        }

        let tail = &self.input[self.position..];
        if contains_ascii_case_insensitive(tail, b"binary") {
            return Ok(ParsedStatement::Unsupported(
                TableDataRepresentation::Binary,
            ));
        }
        if !tail.is_empty() {
            return Ok(ParsedStatement::Unsupported(TableDataRepresentation::Other));
        }

        Err(MetadataParseError::Malformed(
            "COPY statement is missing its semicolon",
        ))
    }

    fn skip_relation(&mut self) -> Result<(), MetadataParseError> {
        self.skip_identifier("COPY relation name is missing")?;
        if self.consume_byte(b'.') {
            self.skip_identifier("qualified COPY relation name is incomplete")?;
            if self.peek_byte() == Some(b'.') {
                return Err(MetadataParseError::Malformed(
                    "COPY relation name has too many qualifiers",
                ));
            }
        }
        Ok(())
    }

    fn parse_columns(&mut self) -> Result<Vec<Column>, MetadataParseError> {
        let mut columns = Vec::new();
        self.skip_whitespace();
        if self.consume_byte(b')') {
            return Ok(columns);
        }

        loop {
            let actual = columns
                .len()
                .checked_add(1)
                .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
            let actual_u64 = to_u64(actual)?;
            if actual > self.max_columns {
                return Err(PgDumpError::CopyColumnCountLimitExceeded {
                    dump_id: self.dump_id.as_i32(),
                    limit: to_u64(self.max_columns)?,
                    actual: actual_u64,
                }
                .into());
            }
            columns.try_reserve(1).map_err(|_| {
                PgDumpError::CopyColumnMetadataAllocationFailed {
                    dump_id: self.dump_id.as_i32(),
                    requested: actual_u64,
                }
            })?;
            columns.push(Column::new(self.parse_column_identifier()?));

            self.skip_whitespace();
            if self.consume_byte(b',') {
                self.skip_whitespace();
                continue;
            }
            if self.consume_byte(b')') {
                return Ok(columns);
            }
            return Err(MetadataParseError::Malformed(
                "COPY column list must use commas and a closing parenthesis",
            ));
        }
    }

    fn parse_column_identifier(&mut self) -> Result<Vec<u8>, MetadataParseError> {
        if self.consume_byte(b'\"') {
            return self.parse_quoted_identifier();
        }

        let start = self.position;
        while let Some(byte) = self.peek_byte() {
            if byte.is_ascii_whitespace() || matches!(byte, b',' | b')') {
                break;
            }
            if matches!(byte, b'(' | b';' | b'\"') {
                return Err(MetadataParseError::Malformed(
                    "invalid byte in unquoted COPY column name",
                ));
            }
            self.position += 1;
        }
        if self.position == start {
            return Err(MetadataParseError::Malformed("COPY column name is missing"));
        }
        clone_bytes(&self.input[start..self.position], self.dump_id).map_err(Into::into)
    }

    fn parse_quoted_identifier(&mut self) -> Result<Vec<u8>, MetadataParseError> {
        let mut logical = Vec::new();
        loop {
            let Some(byte) = self.peek_byte() else {
                return Err(MetadataParseError::Malformed(
                    "quoted COPY column name is unterminated",
                ));
            };
            self.position += 1;

            if byte == b'\"' {
                if self.peek_byte() == Some(b'\"') {
                    self.position += 1;
                    push_byte(&mut logical, b'\"', self.dump_id)?;
                    continue;
                }
                if logical.is_empty() {
                    return Err(MetadataParseError::Malformed(
                        "quoted COPY column name is empty",
                    ));
                }
                return Ok(logical);
            }
            push_byte(&mut logical, byte, self.dump_id)?;
        }
    }

    fn skip_identifier(&mut self, missing: &'static str) -> Result<(), MetadataParseError> {
        if self.consume_byte(b'\"') {
            let mut logical_bytes = 0_usize;
            loop {
                let Some(byte) = self.peek_byte() else {
                    return Err(MetadataParseError::Malformed(
                        "quoted COPY relation name is unterminated",
                    ));
                };
                self.position += 1;
                if byte == b'\"' {
                    if self.peek_byte() == Some(b'\"') {
                        self.position += 1;
                        logical_bytes = logical_bytes
                            .checked_add(1)
                            .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
                        continue;
                    }
                    if logical_bytes == 0 {
                        return Err(MetadataParseError::Malformed(
                            "quoted COPY relation name is empty",
                        ));
                    }
                    return Ok(());
                }
                logical_bytes = logical_bytes
                    .checked_add(1)
                    .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
            }
        }

        let start = self.position;
        while let Some(byte) = self.peek_byte() {
            if byte.is_ascii_whitespace() || matches!(byte, b'.' | b'(') {
                break;
            }
            if matches!(byte, b',' | b')' | b';' | b'\"') {
                return Err(MetadataParseError::Malformed(
                    "invalid byte in COPY relation name",
                ));
            }
            self.position += 1;
        }
        if self.position == start {
            return Err(MetadataParseError::Malformed(missing));
        }
        Ok(())
    }

    fn expect_bytes(
        &mut self,
        expected: &[u8],
        reason: &'static str,
    ) -> Result<(), MetadataParseError> {
        let end = self
            .position
            .checked_add(expected.len())
            .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
        if self.input.get(self.position..end) != Some(expected) {
            return Err(MetadataParseError::Malformed(reason));
        }
        self.position = end;
        Ok(())
    }

    fn require_whitespace(&mut self, reason: &'static str) -> Result<(), MetadataParseError> {
        if !self.skip_whitespace() {
            return Err(MetadataParseError::Malformed(reason));
        }
        Ok(())
    }

    fn skip_whitespace(&mut self) -> bool {
        let start = self.position;
        while self
            .peek_byte()
            .is_some_and(|byte| byte.is_ascii_whitespace())
        {
            self.position += 1;
        }
        self.position != start
    }

    fn consume_byte(&mut self, expected: u8) -> bool {
        if self.peek_byte() == Some(expected) {
            self.position += 1;
            true
        } else {
            false
        }
    }

    fn peek_byte(&self) -> Option<u8> {
        self.input.get(self.position).copied()
    }
}

fn clone_bytes(bytes: &[u8], dump_id: DumpId) -> Result<Vec<u8>, PgDumpError> {
    let requested = to_u64(bytes.len())?;
    let mut copy = Vec::new();
    copy.try_reserve_exact(bytes.len()).map_err(|_| {
        PgDumpError::CopyColumnMetadataAllocationFailed {
            dump_id: dump_id.as_i32(),
            requested,
        }
    })?;
    copy.extend_from_slice(bytes);
    Ok(copy)
}

fn push_byte(bytes: &mut Vec<u8>, byte: u8, dump_id: DumpId) -> Result<(), MetadataParseError> {
    if bytes.len() == bytes.capacity() {
        let requested = bytes
            .len()
            .checked_add(1)
            .ok_or(PgDumpError::ArithmeticOverflow { offset: 0 })?;
        bytes
            .try_reserve(1)
            .map_err(|_| PgDumpError::CopyColumnMetadataAllocationFailed {
                dump_id: dump_id.as_i32(),
                requested: u64::try_from(requested).unwrap_or(u64::MAX),
            })?;
    }
    bytes.push(byte);
    Ok(())
}

fn contains_ascii_case_insensitive(haystack: &[u8], needle: &[u8]) -> bool {
    haystack
        .windows(needle.len())
        .any(|window| window.eq_ignore_ascii_case(needle))
}

fn to_u64(value: usize) -> Result<u64, PgDumpError> {
    u64::try_from(value).map_err(|_| PgDumpError::ArithmeticOverflow { offset: 0 })
}