use std::io;

// Wire format for one atomic batch, encoded incrementally by JS via
// op_kv_atomic_push and committed by op_kv_atomic_commit. All integers are
// little-endian; keyspace names are UTF-8 bytes, keys/values are opaque.
//
//   batch := op*
//   op    := CHECK ks key expected | SET ks key value | DELETE ks key
//   CHECK  := 0x01 u16 ks_len ks u16 key_len key u8 has_expected [u16 stamp_len stamp]
//   SET    := 0x02 u16 ks_len ks u16 key_len key u32 value_len value
//   DELETE := 0x03 u16 ks_len ks u16 key_len key
//
// has_expected = 0 means "must not exist" (null check); 1 carries the exact
// versionstamp the entry must currently hold.

pub const TAG_CHECK: u8 = 1;
pub const TAG_SET: u8 = 2;
pub const TAG_DELETE: u8 = 3;

// Sane staging limits so a hostile or buggy guest cannot grow a batch without
// bound before commit parsing rejects it.
pub const MAX_BATCH_OPS: usize = 1_000;
pub const MAX_BATCH_BYTES: usize = 4 << 20;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BatchOp {
    Check {
        keyspace: String,
        key: Vec<u8>,
        expected: Option<String>,
    },
    Set {
        keyspace: String,
        key: Vec<u8>,
        value: Vec<u8>,
    },
    Delete {
        keyspace: String,
        key: Vec<u8>,
    },
}

impl BatchOp {
    pub fn keyspace(&self) -> &str {
        match self {
            BatchOp::Check { keyspace, .. }
            | BatchOp::Set { keyspace, .. }
            | BatchOp::Delete { keyspace, .. } => keyspace,
        }
    }
}

fn invalid(message: impl Into<String>) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidInput, message.into())
}

struct Cursor<'a> {
    bytes: &'a [u8],
}

impl<'a> Cursor<'a> {
    fn take(&mut self, len: usize) -> io::Result<&'a [u8]> {
        if self.bytes.len() < len {
            return Err(invalid("Truncated atomic batch"));
        }
        let (head, tail) = self.bytes.split_at(len);
        self.bytes = tail;
        Ok(head)
    }

    fn u8(&mut self) -> io::Result<u8> {
        Ok(self.take(1)?[0])
    }

    fn u16(&mut self) -> io::Result<usize> {
        let bytes = self.take(2)?;
        Ok(u16::from_le_bytes([bytes[0], bytes[1]]) as usize)
    }

    fn u32(&mut self) -> io::Result<usize> {
        let bytes = self.take(4)?;
        Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize)
    }

    fn keyspace(&mut self) -> io::Result<String> {
        let len = self.u16()?;
        let bytes = self.take(len)?;
        let name = std::str::from_utf8(bytes).map_err(|_| invalid("Keyspace name is not UTF-8"))?;
        super::store::check_keyspace_name(name)?;
        Ok(name.to_owned())
    }

    fn key(&mut self) -> io::Result<Vec<u8>> {
        let len = self.u16()?;
        Ok(self.take(len)?.to_vec())
    }
}

pub fn decode(bytes: &[u8]) -> io::Result<Vec<BatchOp>> {
    if bytes.len() > MAX_BATCH_BYTES {
        return Err(invalid("Atomic batch exceeds size limit"));
    }
    let mut cursor = Cursor { bytes };
    let mut ops = Vec::new();
    while !cursor.bytes.is_empty() {
        if ops.len() >= MAX_BATCH_OPS {
            return Err(invalid("Atomic batch exceeds operation limit"));
        }
        let keyspace;
        let key;
        match cursor.u8()? {
            TAG_CHECK => {
                keyspace = cursor.keyspace()?;
                key = cursor.key()?;
                let expected = match cursor.u8()? {
                    0 => None,
                    1 => {
                        let len = cursor.u16()?;
                        let stamp = cursor.take(len)?;
                        let stamp = std::str::from_utf8(stamp)
                            .map_err(|_| invalid("Versionstamp is not UTF-8"))?;
                        Some(stamp.to_owned())
                    }
                    _ => return Err(invalid("Invalid check expectation flag")),
                };
                ops.push(BatchOp::Check { keyspace, key, expected });
            }
            TAG_SET => {
                keyspace = cursor.keyspace()?;
                key = cursor.key()?;
                let len = cursor.u32()?;
                let value = cursor.take(len)?.to_vec();
                ops.push(BatchOp::Set { keyspace, key, value });
            }
            TAG_DELETE => {
                keyspace = cursor.keyspace()?;
                key = cursor.key()?;
                ops.push(BatchOp::Delete { keyspace, key });
            }
            tag => return Err(invalid(format!("Unknown atomic batch tag {tag}"))),
        }
    }
    Ok(ops)
}

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

    fn encode(parts: &[&[u8]]) -> Vec<u8> {
        parts.concat()
    }

    #[test]
    fn round_trip() {
        let batch = encode(&[
            &[TAG_CHECK],
            &[2, 0],
            b"ks",
            &[1, 0],
            b"k",
            &[1],
            &[4, 0],
            b"0123",
            &[TAG_SET],
            &[2, 0],
            b"ks",
            &[1, 0],
            b"k",
            &[3, 0, 0, 0],
            b"val",
            &[TAG_DELETE],
            &[1, 0],
            b"q",
            &[0, 0],
        ]);
        assert_eq!(
            decode(&batch).unwrap(),
            vec![
                BatchOp::Check {
                    keyspace: "ks".into(),
                    key: b"k".to_vec(),
                    expected: Some("0123".into()),
                },
                BatchOp::Set {
                    keyspace: "ks".into(),
                    key: b"k".into(),
                    value: b"val".to_vec(),
                },
                BatchOp::Delete {
                    keyspace: "q".into(),
                    key: vec![],
                },
            ]
        );
    }

    #[test]
    fn rejects_truncation_and_unknown_tags() {
        assert!(decode(&[TAG_SET, 2, 0]).is_err());
        assert!(decode(&[9]).is_err());
        assert!(decode(&[]).unwrap().is_empty());
    }
}