use crate::error::Result;
use crate::error::StarError;
use crate::types::{StarCommit, StarItem, StarMstNode};
use crate::validation::validate_node_structure;
use std::io::Write;

pub struct StarEncoder;

impl StarEncoder {
    fn write_varint<W: Write>(val: usize, dst: &mut W) -> std::io::Result<()> {
        let mut buf = unsigned_varint::encode::usize_buffer();
        let encoded = unsigned_varint::encode::usize(val, &mut buf);
        dst.write_all(encoded)
    }

    pub fn write_header<W: Write>(commit: &StarCommit, dst: &mut W) -> Result<()> {
        dst.write_all(&[0x2A])?;

        Self::write_varint(1, dst)?;

        let commit_bytes = serde_ipld_dagcbor::to_vec(commit)
            .map_err(|e| crate::error::StarError::Cbor(e.to_string()))?;

        Self::write_varint(commit_bytes.len(), dst)?;
        dst.write_all(&commit_bytes)?;

        Ok(())
    }

    pub fn write_node<W: Write>(node: &StarMstNode, dst: &mut W) -> Result<()> {
        let node_bytes = serde_ipld_dagcbor::to_vec(node)
            .map_err(|e| crate::error::StarError::Cbor(e.to_string()))?;

        Self::write_varint(node_bytes.len(), dst)?;
        dst.write_all(&node_bytes)?;

        Ok(())
    }

    pub fn write_record<W: Write>(record_bytes: &[u8], dst: &mut W) -> Result<()> {
        Self::write_varint(record_bytes.len(), dst)?;
        dst.write_all(record_bytes)?;

        Ok(())
    }
}

#[cfg(feature = "async")]
impl tokio_util::codec::Encoder<StarItem> for StarEncoder {
    type Error = crate::error::StarError;

    fn encode(&mut self, item: StarItem, dst: &mut bytes::BytesMut) -> Result<()> {
        use bytes::BufMut;
        // BytesMut::writer() returns an impl Write
        let mut writer = dst.writer();

        match item {
            StarItem::Commit(c) => Self::write_header(&c, &mut writer),
            StarItem::Node(n) => Self::write_node(&n, &mut writer),
            StarItem::Record { content, .. } => {
                if let Some(bytes) = content {
                    Self::write_record(&bytes, &mut writer)
                } else {
                    Err(crate::error::StarError::InvalidState(
                        "Cannot serialize record without content".into(),
                    ))
                }
            }
        }
    }
}

/// A serializer that enforces strict STAR format compliance.
pub struct StarSerializer<W> {
    writer: W,
    validator: StarValidator,
}

impl<W: Write> StarSerializer<W> {
    pub fn new(writer: W) -> Self {
        Self {
            writer,
            validator: StarValidator::new(),
        }
    }

    pub fn write_header(&mut self, commit: &StarCommit) -> Result<()> {
        self.validator.accept_header(commit)?;
        StarEncoder::write_header(commit, &mut self.writer)
    }

    pub fn write_node(&mut self, node: &StarMstNode) -> Result<()> {
        self.validator.accept_node(node)?;
        StarEncoder::write_node(node, &mut self.writer)
    }

    pub fn write_record(&mut self, record_bytes: &[u8]) -> Result<()> {
        self.validator.accept_record(record_bytes)?;
        StarEncoder::write_record(record_bytes, &mut self.writer)
    }

    pub fn finish(self) -> Result<W> {
        if !self.validator.is_done() {
            return Err(crate::error::StarError::InvalidState(
                "Incomplete tree".into(),
            ));
        }
        Ok(self.writer)
    }
}

// Validator State Machine (Moved from validation.rs)

#[derive(Debug, Default)]
pub struct StarValidator {
    state: ValidatorState,
}

#[derive(Debug, Default)]
enum ValidatorState {
    #[default]
    Header,
    Body {
        stack: Vec<Expectation>,
    },
}

#[derive(Debug)]
enum Expectation {
    Root,
    Node { height: u32 },
    Record,
}

impl StarValidator {
    pub fn new() -> Self {
        Self {
            state: ValidatorState::Header,
        }
    }

    pub fn accept_header(&mut self, commit: &StarCommit) -> Result<()> {
        match &self.state {
            ValidatorState::Header => {
                let stack = if commit.data.is_some() {
                    vec![Expectation::Root]
                } else {
                    Vec::new() // Empty tree
                };
                self.state = ValidatorState::Body { stack };
                Ok(())
            }
            _ => Err(StarError::InvalidState(
                "Header already written or invalid state".into(),
            )),
        }
    }

    pub fn accept_node(&mut self, node: &StarMstNode) -> Result<()> {
        match &mut self.state {
            ValidatorState::Body { stack } => {
                if stack.is_empty() {
                    return Err(StarError::InvalidState(
                        "Unexpected node: tree is complete".into(),
                    ));
                }

                let expectation = stack.pop().unwrap();
                let expected_height = match expectation {
                    Expectation::Record => {
                        return Err(StarError::InvalidState("Expected record, got node".into()));
                    }
                    Expectation::Root => None,
                    Expectation::Node { height } => Some(height),
                };

                // Use the shared validation logic
                let (height, _) = validate_node_structure(node, expected_height)?;

                let child_height = if height > 0 { height - 1 } else { 0 };

                for e in node.e.iter().rev() {
                    if e.t_archived == Some(true) {
                        stack.push(Expectation::Node {
                            height: child_height,
                        });
                    }
                    if e.v_archived == Some(true) {
                        stack.push(Expectation::Record);
                    }
                }

                if node.l_archived == Some(true) {
                    stack.push(Expectation::Node {
                        height: child_height,
                    });
                }

                Ok(())
            }
            _ => Err(StarError::InvalidState("Invalid state for node".into())),
        }
    }

    pub fn accept_record(&mut self, _bytes: &[u8]) -> Result<()> {
        match &mut self.state {
            ValidatorState::Body { stack } => {
                if stack.is_empty() {
                    return Err(StarError::InvalidState(
                        "Unexpected record: tree is complete".into(),
                    ));
                }
                match stack.pop().unwrap() {
                    Expectation::Record => Ok(()),
                    _ => Err(StarError::InvalidState("Expected node, got record".into())),
                }
            }
            _ => Err(StarError::InvalidState("Invalid state for record".into())),
        }
    }

    pub fn is_done(&self) -> bool {
        match &self.state {
            ValidatorState::Body { stack } => stack.is_empty(),
            _ => false,
        }
    }
}