use bytes::{Buf, BytesMut};
use std::collections::HashMap;

/// PostgreSQL logical replication message parser
/// Handles pgoutput protocol messages from replication streams

#[derive(Debug, Clone)]
pub struct ColumnInfo {
    pub flags: u8,              // 1 if part of key, 0 otherwise
    pub name: String,
    pub type_oid: u32,
    pub type_modifier: i32,
}

#[derive(Debug, Clone)]
pub struct RelationInfo {
    pub oid: u32,
    pub namespace: String,
    pub name: String,
    pub replica_identity: u8,
    pub columns: Vec<ColumnInfo>,
}

#[derive(Debug, Clone)]
pub enum ColumnValue {
    Null,
    UnchangedToasted,
    Text(String),
    Binary(Vec<u8>),
}

#[derive(Debug)]
pub struct TupleData {
    pub values: Vec<ColumnValue>,
}

#[derive(Debug)]
pub enum LogicalMessage {
    Begin {
        final_lsn: u64,
        timestamp: i64,
        xid: u32,
    },
    Relation(RelationInfo),
    Insert {
        xid: u32,
        rel_oid: u32,
        new_tuple: TupleData,
    },
    Update {
        xid: u32,
        rel_oid: u32,
        key_tuple: Option<TupleData>,
        old_tuple: Option<TupleData>,
        new_tuple: TupleData,
    },
    Delete {
        xid: u32,
        rel_oid: u32,
        key_tuple: Option<TupleData>,
        old_tuple: Option<TupleData>,
    },
    Commit {
        flags: u8,
        commit_lsn: u64,
        end_lsn: u64,
        timestamp: i64,
    },
    Type {
        xid: u32,
        type_oid: u32,
        namespace: String,
        type_name: String,
    },
    Truncate {
        xid: u32,
        num_relations: u32,
        options: u8,
        rel_oids: Vec<u32>,
    },
    Message {
        xid: u32,
        flags: u8,
        lsn: u64,
        prefix: String,
        content: Vec<u8>,
    },
}

pub struct MessageParser {
    /// Map of relation OID -> RelationInfo for decoding tuples
    relations: HashMap<u32, RelationInfo>,
    /// Map of type OID -> type name for potential future use
    types: HashMap<u32, (String, String)>,
}

impl MessageParser {
    pub fn new() -> Self {
        Self {
            relations: HashMap::new(),
            types: HashMap::new(),
        }
    }

    /// Parse a single logical replication message from bytes
    /// Assumes input starts with the message type byte
    pub fn parse_message(&mut self, input: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        if input.is_empty() {
            return Ok(None);
        }

        let msg_type = input[0] as char;

        match msg_type {
            'B' => self.parse_begin(input),
            'I' => self.parse_insert(input),
            'U' => self.parse_update(input),
            'D' => self.parse_delete(input),
            'C' => self.parse_commit(input),
            'R' => {
                self.parse_relation(input)?;
                Ok(None) // Relation messages are stored internally
            }
            'Y' => self.parse_type(input),
            'T' => self.parse_truncate(input),
            'M' => self.parse_message_msg(input),
            't' => Err("Tuple data outside of Insert/Update/Delete (unexpected)".to_string()),
            _ => Err(format!("Unknown message type: {} (byte: {})", msg_type, input[0])),
        }
    }

    fn read_string(buf: &mut BytesMut) -> Result<String, String> {
        // PostgreSQL strings are null-terminated
        let mut vec = Vec::new();
        loop {
            if buf.is_empty() {
                return Err("EOF while reading string".to_string());
            }
            let byte = buf.get_u8();
            if byte == 0 {
                break;
            }
            vec.push(byte);
        }
        String::from_utf8(vec).map_err(|e| format!("Invalid UTF-8 in string: {}", e))
    }

    fn read_i64(buf: &mut BytesMut) -> Result<i64, String> {
        if buf.len() < 8 {
            return Err("Not enough bytes for i64".to_string());
        }
        Ok(buf.get_i64())
    }

    fn read_u64(buf: &mut BytesMut) -> Result<u64, String> {
        if buf.len() < 8 {
            return Err("Not enough bytes for u64".to_string());
        }
        Ok(buf.get_u64())
    }

    fn read_i32(buf: &mut BytesMut) -> Result<i32, String> {
        if buf.len() < 4 {
            return Err("Not enough bytes for i32".to_string());
        }
        Ok(buf.get_i32())
    }

    fn read_u32(buf: &mut BytesMut) -> Result<u32, String> {
        if buf.len() < 4 {
            return Err("Not enough bytes for u32".to_string());
        }
        Ok(buf.get_u32())
    }

    fn read_i16(buf: &mut BytesMut) -> Result<i16, String> {
        if buf.len() < 2 {
            return Err("Not enough bytes for i16".to_string());
        }
        Ok(buf.get_i16())
    }

    fn read_u8(buf: &mut BytesMut) -> Result<u8, String> {
        if buf.is_empty() {
            return Err("Not enough bytes for u8".to_string());
        }
        Ok(buf.get_u8())
    }

    fn parse_begin(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'B'
        let final_lsn = Self::read_u64(buf)?;
        let timestamp = Self::read_i64(buf)?;
        let xid = Self::read_u32(buf)?;

        Ok(Some(LogicalMessage::Begin {
            final_lsn,
            timestamp,
            xid,
        }))
    }

    fn parse_commit(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'C'
        let flags = Self::read_u8(buf)?;
        let commit_lsn = Self::read_u64(buf)?;
        let end_lsn = Self::read_u64(buf)?;
        let timestamp = Self::read_i64(buf)?;

        Ok(Some(LogicalMessage::Commit {
            flags,
            commit_lsn,
            end_lsn,
            timestamp,
        }))
    }

    fn parse_relation(&mut self, buf: &mut BytesMut) -> Result<(), String> {
        buf.get_u8(); // 'R'
        // xid only present for streamed transactions (version 2+), skip for now
        let rel_oid = Self::read_u32(buf)?;
        let namespace = Self::read_string(buf)?;
        let name = Self::read_string(buf)?;
        let replica_identity = Self::read_u8(buf)?;
        let num_columns = Self::read_i16(buf)? as usize;

        let mut columns = Vec::with_capacity(num_columns);
        for _ in 0..num_columns {
            let flags = Self::read_u8(buf)?;
            let col_name = Self::read_string(buf)?;
            let type_oid = Self::read_u32(buf)?;
            let type_modifier = Self::read_i32(buf)?;

            columns.push(ColumnInfo {
                flags,
                name: col_name,
                type_oid,
                type_modifier,
            });
        }

        let relation = RelationInfo {
            oid: rel_oid,
            namespace,
            name,
            replica_identity,
            columns,
        };

        self.relations.insert(rel_oid, relation);
        Ok(())
    }

    fn parse_insert(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'I'
        let xid = Self::read_u32(buf)?;
        let rel_oid = Self::read_u32(buf)?;

        // Next byte must be 'N' for new tuple
        let marker = Self::read_u8(buf)?;
        if marker != b'N' {
            return Err(format!("Expected 'N' in INSERT, got: {}", marker as char));
        }

        let new_tuple = self.parse_tuple_data(buf, rel_oid)?;

        Ok(Some(LogicalMessage::Insert {
            xid,
            rel_oid,
            new_tuple,
        }))
    }

    fn parse_update(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'U'
        let xid = Self::read_u32(buf)?;
        let rel_oid = Self::read_u32(buf)?;

        let mut key_tuple = None;
        let mut old_tuple = None;

        // Optional K or O markers
        if !buf.is_empty() {
            match buf[0] as char {
                'K' => {
                    buf.get_u8(); // consume 'K'
                    key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
                }
                'O' => {
                    buf.get_u8(); // consume 'O'
                    old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
                }
                _ => {} // neither K nor O, proceed to N
            }
        }

        let marker = Self::read_u8(buf)?;
        if marker != b'N' {
            return Err(format!("Expected 'N' in UPDATE, got: {}", marker as char));
        }

        let new_tuple = self.parse_tuple_data(buf, rel_oid)?;

        Ok(Some(LogicalMessage::Update {
            xid,
            rel_oid,
            key_tuple,
            old_tuple,
            new_tuple,
        }))
    }

    fn parse_delete(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'D'
        let xid = Self::read_u32(buf)?;
        let rel_oid = Self::read_u32(buf)?;

        let mut key_tuple = None;
        let mut old_tuple = None;

        // Either K or O (but not both)
        if !buf.is_empty() {
            match buf[0] as char {
                'K' => {
                    buf.get_u8(); // consume 'K'
                    key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
                }
                'O' => {
                    buf.get_u8(); // consume 'O'
                    old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
                }
                _ => return Err(format!("Expected K or O in DELETE, got: {}", buf[0] as char)),
            }
        }

        Ok(Some(LogicalMessage::Delete {
            xid,
            rel_oid,
            key_tuple,
            old_tuple,
        }))
    }

    fn parse_type(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'Y'
        let type_oid = Self::read_u32(buf)?;
        let namespace = Self::read_string(buf)?;
        let type_name = Self::read_string(buf)?;
        let xid = 0; // Type doesn't have xid in base protocol

        self.types
            .insert(type_oid, (namespace.clone(), type_name.clone()));

        Ok(Some(LogicalMessage::Type {
            xid,
            type_oid,
            namespace,
            type_name,
        }))
    }

    fn parse_truncate(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'T'
        let xid = Self::read_u32(buf)?;
        let num_relations = Self::read_u32(buf)?;
        let options = Self::read_u8(buf)?;

        let mut rel_oids = Vec::with_capacity(num_relations as usize);
        for _ in 0..num_relations {
            rel_oids.push(Self::read_u32(buf)?);
        }

        Ok(Some(LogicalMessage::Truncate {
            xid,
            num_relations,
            options,
            rel_oids,
        }))
    }

    fn parse_message_msg(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
        buf.get_u8(); // 'M'
        let xid = Self::read_u32(buf)?;
        let flags = Self::read_u8(buf)?;
        let lsn = Self::read_u64(buf)?;
        let prefix = Self::read_string(buf)?;
        let content_len = Self::read_i32(buf)? as usize;

        if buf.len() < content_len {
            return Err(format!(
                "Not enough bytes for message content: {} < {}",
                buf.len(),
                content_len
            ));
        }

        let content = buf.split_to(content_len).to_vec();

        Ok(Some(LogicalMessage::Message {
            xid,
            flags,
            lsn,
            prefix,
            content,
        }))
    }

    fn parse_tuple_data(&self, buf: &mut BytesMut, rel_oid: u32) -> Result<TupleData, String> {
        let num_columns = Self::read_i16(buf)? as usize;
        let mut values = Vec::with_capacity(num_columns);

        let relation = self
            .relations
            .get(&rel_oid)
            .ok_or_else(|| format!("Unknown relation OID: {}", rel_oid))?;

        if num_columns != relation.columns.len() {
            return Err(format!(
                "Column count mismatch: {} vs {}",
                num_columns,
                relation.columns.len()
            ));
        }

        for _ in 0..num_columns {
            let marker = Self::read_u8(buf)?;

            let value = match marker as char {
                'n' => ColumnValue::Null,
                'u' => ColumnValue::UnchangedToasted,
                't' => {
                    let len = Self::read_i32(buf)? as usize;
                    if buf.len() < len {
                        return Err("Not enough bytes for text column".to_string());
                    }
                    let data = buf.split_to(len).to_vec();
                    let text = String::from_utf8(data)
                        .map_err(|e| format!("Invalid UTF-8 in column: {}", e))?;
                    ColumnValue::Text(text)
                }
                'b' => {
                    let len = Self::read_i32(buf)? as usize;
                    if buf.len() < len {
                        return Err("Not enough bytes for binary column".to_string());
                    }
                    let data = buf.split_to(len).to_vec();
                    ColumnValue::Binary(data)
                }
                _ => return Err(format!("Unknown column marker: {}", marker as char)),
            };

            values.push(value);
        }

        Ok(TupleData { values })
    }

    /// Get a relation by OID (used for cache invalidation mapping)
    pub fn get_relation(&self, oid: u32) -> Option<&RelationInfo> {
        self.relations.get(&oid)
    }

    /// Get all cached relations
    pub fn relations(&self) -> &HashMap<u32, RelationInfo> {
        &self.relations
    }
}

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

    #[test]
    fn test_message_parser_creation() {
        let parser = MessageParser::new();
        assert!(parser.relations.is_empty());
    }
}