use bytes::Buf as _;

use super::number::NumberDecoder;
use crate::keyspace::NULLSPACEID;

pub const TXN_KEY_PREFIX: u8 = b'x';
const TIDB_META_KEY_PREFIX: u8 = b'm';
const TIDB_TABLE_KEY_PREFIX: u8 = b't';
const DEFAULT_KEY_SPACE_ID: [u8; 3] = [0, 0, 0]; // reserve 3 bytes for key space id.
const KEYSPACE_ID_LEN: usize = DEFAULT_KEY_SPACE_ID.len();
pub const KEYSPACE_PREFIX_LEN: usize = KEYSPACE_ID_LEN + 1;

/// The key mode inferred from the key prefix.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum KeyMode {
    /// Raw key.
    Raw,
    /// Transaction key.
    Txn,
    /// TiDB key.
    ///
    /// It doesn't mean that the key is certainly written by
    /// TiDB, but instead, it means that the key matches the definition of
    /// TiDB key in API V2, therefore, the key is treated as TiDB data in
    /// order to fulfill compatibility.
    Tidb,
    /// Unrecognised key mode.
    Unknown,
}

fn parse_key_mode(key: &[u8]) -> KeyMode {
    if key.len() < KEYSPACE_PREFIX_LEN {
        return KeyMode::Unknown;
    }

    match key[0] {
        TXN_KEY_PREFIX => KeyMode::Txn,
        TIDB_META_KEY_PREFIX | TIDB_TABLE_KEY_PREFIX => KeyMode::Tidb,
        _ => KeyMode::Unknown,
    }
}

/// Return `None` when the key is not an API V2 key.
pub fn get_u32_keyspace_id_by_key(key: &[u8]) -> Option<u32> {
    let key_mode = parse_key_mode(key);
    if key_mode == KeyMode::Txn {
        Some([0, key[1], key[2], key[3]].as_slice().get_u32())
    } else {
        None
    }
}

pub fn get_keyspace_prefix(key: &[u8]) -> Option<&[u8]> {
    matches!(parse_key_mode(key), KeyMode::Txn).then(|| &key[0..KEYSPACE_PREFIX_LEN])
}

pub fn get_keyspace_table_id(key: &[u8]) -> Option<(u32, i64)> {
    let keyspace_id = get_u32_keyspace_id_by_key(key)?;
    let mut buf = &key[KEYSPACE_PREFIX_LEN..];
    if buf.is_empty() || buf[0] != b't' {
        return None;
    }
    buf = &buf[1..];
    let table_id = buf.read_i64().ok()?;
    Some((keyspace_id, table_id))
}

pub fn get_txn_keyspace_prefix(keyspace_id: u32) -> Vec<u8> {
    let mut start_key = keyspace_id.to_be_bytes();
    start_key[0] = TXN_KEY_PREFIX;
    start_key.to_vec()
}

pub fn add_keyspace_prefix(keyspace_id: u32, handle: &[u8]) -> Vec<u8> {
    if keyspace_id == NULLSPACEID {
        handle.to_vec()
    } else {
        let mut v = Vec::with_capacity(handle.len() + KEYSPACE_PREFIX_LEN);
        v.extend_from_slice(get_txn_keyspace_prefix(keyspace_id).as_slice());
        v.extend_from_slice(handle);
        v
    }
}

pub fn remove_keyspace_prefix(handle: &[u8]) -> &[u8] {
    if get_u32_keyspace_id_by_key(handle).is_some() {
        &handle[KEYSPACE_PREFIX_LEN..]
    } else {
        handle
    }
}