// Copyright 2023 PingCAP, Inc.

use std::sync::Arc;

use async_trait::async_trait;
use tikv_util::{
    codec::{
        bytes::{decode_bytes, encode_bytes},
        number::NumberEncoder,
    },
    debug, error, warn,
};
use txn_types::TimeStamp;

use crate::schema::{DbInfo, STATE_PUBLIC, TableInfo};

const M_DBS: &[u8] = b"mDBs";
const M_TABLE_PREFIX: &[u8] = b"Table";
pub const DEFAULT_SCAN_BATCH_SIZE: u32 = 1024;

pub async fn load_schema(
    kv_scanner: Arc<dyn KvScanner>,
    keyspace: &[u8],
    start_ts: TimeStamp,
) -> Result<Vec<DbInfo>, String> {
    load_schema_internal(kv_scanner, keyspace, start_ts, false).await
}

pub(crate) async fn load_schema_with_required_changes(
    kv_scanner: Arc<dyn KvScanner>,
    keyspace: &[u8],
    start_ts: TimeStamp,
) -> Result<Vec<DbInfo>, String> {
    load_schema_internal(kv_scanner, keyspace, start_ts, true).await
}

async fn load_schema_internal(
    kv_scanner: Arc<dyn KvScanner>,
    keyspace: &[u8],
    start_ts: TimeStamp,
    only_required_changes: bool,
) -> Result<Vec<DbInfo>, String> {
    let mut dbs = vec![];
    let (mut start, end) = dbs_range_keys(keyspace);
    loop {
        let pairs = kv_scanner.scan(&start, &end, start_ts).await?;
        if pairs.is_empty() {
            break;
        }
        debug!(
            "scan db batch pairs {}, start {:?}, end {:?}",
            pairs.len(),
            start,
            end
        );
        start = next_scan_start_key(&pairs);

        for (key, value) in pairs {
            let mut db = parse_db_info(&key, &value)?;
            if db.state != STATE_PUBLIC {
                continue;
            }
            db.tables = load_db_tables(
                kv_scanner.as_ref(),
                keyspace,
                db.id,
                start_ts,
                only_required_changes,
            )
            .await?;
            if !only_required_changes || !db.tables.is_empty() {
                dbs.push(db);
            }
        }
    }
    Ok(dbs)
}

async fn load_db_tables(
    kv_scanner: &dyn KvScanner,
    keyspace: &[u8],
    db_id: i64,
    start_ts: TimeStamp,
    only_required_changes: bool,
) -> Result<Vec<TableInfo>, String> {
    let start_key = db_tables_start_key(keyspace, db_id);
    let prefix_len = start_key.len();
    let (mut start, end) = get_range_keys(start_key);
    let mut tables = vec![];

    loop {
        let pairs = kv_scanner.scan(&start, &end, start_ts).await?;
        if pairs.is_empty() {
            break;
        }
        start = next_scan_start_key(&pairs);

        for (key, value) in pairs {
            let mut suffix = &key[prefix_len..];
            let field = decode_bytes(&mut suffix, false).unwrap();
            if !field.starts_with(M_TABLE_PREFIX) {
                continue;
            }
            let res: serde_json::Result<TableInfo> = serde_json::from_slice(&value);
            match res {
                Ok(tbl) => {
                    if tbl.state == STATE_PUBLIC
                        && (!only_required_changes || tbl.with_required_changes())
                    {
                        tables.push(tbl);
                    }
                }
                Err(err) => {
                    let val_str = String::from_utf8_lossy(&value);
                    warn!("invalid key {:?}, val {} err {:?}", key, val_str, err);
                }
            }
        }
    }

    Ok(tables)
}

fn parse_db_info(key: &[u8], value: &[u8]) -> Result<DbInfo, String> {
    serde_json::from_slice(value).map_err(|err| {
        let key = log_wrappers::hex_encode_upper(key);
        let val = String::from_utf8_lossy(value);
        error!(
            "load schema: parse db info failed";
            "err" => ?err, "key" => &key, "val" => val.as_ref(),
        );
        format!("parse db info failed for {}", key)
    })
}

fn dbs_range_keys(keyspace: &[u8]) -> (Vec<u8>, Vec<u8>) {
    let mut start = Vec::with_capacity(keyspace.len() + M_DBS.len());
    start.extend_from_slice(keyspace);
    start.extend_from_slice(M_DBS);
    get_range_keys(start)
}

fn db_tables_start_key(keyspace: &[u8], db_id: i64) -> Vec<u8> {
    let mut db_tables_key = vec![];
    db_tables_key.extend_from_slice(keyspace);
    db_tables_key.push(b'm');
    let encoded_db_str = encode_bytes(format!("DB:{}", db_id).as_bytes());
    db_tables_key.extend_from_slice(encoded_db_str.as_slice());
    db_tables_key.encode_u64('h' as u64).unwrap();
    db_tables_key
}

fn get_range_keys(start_key: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
    let mut end_key = start_key.to_vec();
    let last_idx = end_key.len() - 1;
    assert!(end_key[last_idx] != 0xff);
    end_key[last_idx] += 1;
    (start_key, end_key)
}

fn next_scan_start_key(pairs: &[(Vec<u8>, Vec<u8>)]) -> Vec<u8> {
    let mut next = pairs.last().unwrap().0.clone();
    next.push(0);
    next
}

#[async_trait]
pub trait KvScanner: Send + Sync {
    /// Scans one batch of kv pairs from `[start, end)`.
    ///
    /// Callers should continue scanning from the last returned key plus a
    /// trailing `0` byte until an empty batch is returned. Implementations may
    /// return the whole remaining range when their configured batch size is
    /// `0`.
    async fn scan(
        &self,
        start: &[u8],
        end: &[u8],
        start_ts: TimeStamp,
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, String>;
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, sync::Arc};

    use api_version::ApiV2;
    use futures::executor::block_on;

    use super::*;
    use crate::schema::{CiStr, ColumnInfo, DbInfo, FieldType};

    const TIDB_DBS: &[u8] = b"DBs";

    #[derive(Default)]
    struct MockKvScanner {
        data: HashMap<Vec<u8>, Vec<u8>>,
        scan_batch_size: usize,
    }

    #[async_trait]
    impl KvScanner for MockKvScanner {
        async fn scan(
            &self,
            start: &[u8],
            end: &[u8],
            _start_ts: TimeStamp,
        ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, String> {
            let batch_size = if self.scan_batch_size == 0 {
                usize::MAX
            } else {
                self.scan_batch_size
            };
            let mut pairs = self
                .data
                .iter()
                .filter(|(key, _)| key.as_slice() >= start && key.as_slice() < end)
                .map(|(key, value)| (key.clone(), value.clone()))
                .collect::<Vec<_>>();
            pairs.sort_by(|a, b| a.0.cmp(&b.0));
            pairs.truncate(batch_size);
            Ok(pairs)
        }
    }

    #[test]
    fn test_load_schema_reads_all_tables_across_batches() {
        let keyspace_id = 51;
        let db_id = 163;
        let mut kv = MockKvScanner {
            scan_batch_size: 1,
            ..Default::default()
        };
        kv.data.insert(
            dbs_db_key(keyspace_id, db_id),
            serde_json::to_vec(&public_db(db_id)).unwrap(),
        );
        kv.data.insert(
            schema_data_key(keyspace_id, db_id, 501),
            serde_json::to_vec(&public_table(501)).unwrap(),
        );
        kv.data.insert(
            schema_data_key(keyspace_id, db_id, 502),
            serde_json::to_vec(&public_table(502)).unwrap(),
        );

        let keyspace = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let dbs = block_on(load_schema(Arc::new(kv), &keyspace, TimeStamp::default())).unwrap();
        assert_eq!(dbs.len(), 1);
        assert_eq!(dbs[0].id, db_id);
        let mut table_ids = dbs[0].tables.iter().map(|tbl| tbl.id).collect::<Vec<_>>();
        table_ids.sort_unstable();
        assert_eq!(table_ids, vec![501, 502]);
    }

    #[test]
    fn test_load_schema_with_required_changes_filters_tables_across_batches() {
        let keyspace_id = 52;
        let mut kv = MockKvScanner {
            scan_batch_size: 1,
            ..Default::default()
        };
        kv.data.insert(
            dbs_db_key(keyspace_id, 163),
            serde_json::to_vec(&public_db(163)).unwrap(),
        );
        kv.data.insert(
            dbs_db_key(keyspace_id, 164),
            serde_json::to_vec(&public_db(164)).unwrap(),
        );
        kv.data.insert(
            schema_data_key(keyspace_id, 163, 501),
            serde_json::to_vec(&public_table(501)).unwrap(),
        );
        kv.data.insert(
            schema_data_key(keyspace_id, 163, 502),
            serde_json::to_vec(&public_columnar_table(502)).unwrap(),
        );
        kv.data.insert(
            schema_data_key(keyspace_id, 164, 601),
            serde_json::to_vec(&public_table(601)).unwrap(),
        );

        let keyspace = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let dbs = block_on(load_schema_with_required_changes(
            Arc::new(kv),
            &keyspace,
            TimeStamp::default(),
        ))
        .unwrap();

        assert_eq!(dbs.len(), 1);
        assert_eq!(dbs[0].id, 163);
        assert_eq!(dbs[0].tables.len(), 1);
        assert_eq!(dbs[0].tables[0].id, 502);
    }

    fn public_db(id: i64) -> DbInfo {
        DbInfo {
            id,
            name: CiStr {
                o: format!("db_{id}"),
                l: format!("db_{id}"),
            },
            state: STATE_PUBLIC,
            ..Default::default()
        }
    }

    fn public_table(id: i64) -> TableInfo {
        TableInfo {
            id,
            name: CiStr {
                o: format!("t_{id}"),
                l: format!("t_{id}"),
            },
            state: STATE_PUBLIC,
            ..Default::default()
        }
    }

    fn public_columnar_table(id: i64) -> TableInfo {
        TableInfo {
            cols: Some(vec![ColumnInfo {
                id: 1,
                offset: 0,
                state: STATE_PUBLIC,
                field_type: FieldType {
                    tp: 3,
                    charset: "binary".to_string(),
                    collate: "binary".to_string(),
                    ..Default::default()
                },
                ..Default::default()
            }]),
            comment: "columnar_engine".to_string(),
            ..public_table(id)
        }
    }

    fn dbs_db_key(keyspace_id: u32, db_id: i64) -> Vec<u8> {
        let mut key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let enc_dbs = encode_bytes(TIDB_DBS);
        let enc_db_key = encode_bytes(format!("DB:{db_id}").as_bytes());
        key.push(b'm');
        key.extend_from_slice(&enc_dbs);
        key.encode_u64(b'h' as u64).unwrap();
        key.extend_from_slice(&enc_db_key);
        key
    }

    fn schema_data_key(keyspace_id: u32, db_id: i64, table_id: i64) -> Vec<u8> {
        let mut key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let enc_db_key = encode_bytes(format!("DB:{db_id}").as_bytes());
        let enc_table_key = encode_bytes(format!("Table:{table_id}").as_bytes());
        key.push(b'm');
        key.extend_from_slice(&enc_db_key);
        key.encode_u64(b'h' as u64).unwrap();
        key.extend_from_slice(&enc_table_key);
        key
    }
}