use std::collections::BTreeMap;
use std::io;

use super::{Decision, Store, check_keyspace_name, parse_stamp};

// Bounds on one batch so a hostile or buggy guest cannot make a single
// commit arbitrarily large. Bytes count keyspace names, keys and values.
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: &str) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidInput, message)
}

impl Store {
    pub fn atomic_commit(&self, ops: &[BatchOp]) -> io::Result<(bool, Option<String>)> {
        if ops.len() > MAX_BATCH_OPS {
            return Err(invalid("Atomic batch exceeds operation limit"));
        }
        let mut bytes = 0;
        for op in ops {
            check_keyspace_name(op.keyspace())?;
            bytes += op.keyspace().len()
                + match op {
                    BatchOp::Check { key, .. } | BatchOp::Delete { key, .. } => key.len(),
                    BatchOp::Set { key, value, .. } => key.len() + value.len(),
                };
            // Malformed stamps are caller error, not a quiet check failure.
            if let BatchOp::Check { expected: Some(stamp), .. } = op {
                parse_stamp(stamp)?;
            }
        }
        if bytes > MAX_BATCH_BYTES {
            return Err(invalid("Atomic batch exceeds size limit"));
        }
        if ops.is_empty() {
            return Ok((true, None));
        }
        self.write_tx(|tx| {
            // Checks observe pre-commit state and never register keyspaces: a
            // missing registry row means the entry is absent.
            for op in ops {
                if let BatchOp::Check { keyspace, key, expected } = op {
                    let current = tx.get(keyspace, key.as_slice())?;
                    let pass = match (&current, expected) {
                        (None, None) => true,
                        (Some(_), None) => false,
                        (None, Some(_)) => false,
                        // Both sides are canonical renderings (parse_stamp
                        // accepted `want` above), so string equality suffices.
                        (Some((_, have)), Some(want)) => have == want,
                    };
                    if !pass {
                        // Abort drops the transaction uncommitted: a failed
                        // batch changes nothing, not even metadata.
                        return Ok(Decision::Abort((false, None)));
                    }
                }
            }
            let mut mutations: BTreeMap<(&str, &[u8]), Option<&[u8]>> = BTreeMap::new();
            for op in ops {
                match op {
                    BatchOp::Check { .. } => {}
                    BatchOp::Set { keyspace, key, value } => {
                        mutations
                            .insert((keyspace.as_str(), key.as_slice()), Some(value.as_slice()));
                    }
                    BatchOp::Delete { keyspace, key } => {
                        mutations.insert((keyspace.as_str(), key.as_slice()), None);
                    }
                }
            }
            // A check-only batch commits nothing and burns no versionstamp.
            if mutations.is_empty() {
                return Ok(Decision::Abort((true, None)));
            }
            for ((keyspace, key), mutation) in &mutations {
                match mutation {
                    Some(value) => {
                        tx.put(keyspace, key, value)?;
                    }
                    None => {
                        tx.delete(keyspace, key)?;
                    }
                }
            }
            let stamp = match tx.stamp() {
                Some(stamp) => stamp,
                // Every mutation missed: the batch still consumes a commit.
                None => tx.allocate_stamp()?,
            };
            Ok(Decision::Commit((true, Some(stamp))))
        })
    }
}