// Copyright 2024 PingCAP, Inc.

use std::{collections::HashMap, convert::TryFrom, sync::Arc};

use bytes::{Buf, Bytes};
use cloud_encryption::EncryptionKey;
use http::{Response, StatusCode, header, request::Parts};
use hyper::Body;
use kvengine::{
    dfs::{self, Dfs},
    table::{
        ChecksumType, InnerKey, PessimisticAction, TXN_FILE_FORMAT_V1, TXN_FILE_FORMAT_V2,
        txn_file::TxnChunkBuilder,
    },
};

use crate::{
    common::{get_param, make_response},
    encryption,
    error::Result,
    server::Context,
};
pub(crate) const TARGET_BLOCK_SIZE_DEF: usize = 65536; // 64KB

// the initial API version. Use the same value as `TXN_FILE_FORMAT` value.
const TXN_CHUNK_API_V1: u32 = 1;
// version 2, change the entry format from (op, key, value) to (op,
// pessimistic_action(1byte), key, value); the extra 1 byte "pessimistic_action"
// is used for pessimistic txn mutation.
const TXN_CHUNK_API_V2: u32 = 2;

pub fn is_valid_txn_chunk_api_version(v: u32) -> bool {
    matches!(v, TXN_CHUNK_API_V1 | TXN_CHUNK_API_V2)
}

fn api_version_to_txn_file_format_version(v: u32) -> u16 {
    match v {
        TXN_CHUNK_API_V1 => TXN_FILE_FORMAT_V1,
        TXN_CHUNK_API_V2 => TXN_FILE_FORMAT_V2,
        _ => unreachable!(),
    }
}

/// Txn Chunk API
///
/// * GET /txn_chunk?keyspace_id=<keyspace_id>: query txn file is available for
///   this keyspace or not
///
///   Return: { "available": true/false }
///
/// * POST /txn_chunk?keyspace_id=<keyspace_id>&api_version=<version>: submit
///   data for creating txn chunk.
///
///   Return: { "chunk_id": <chunk_id> }
pub(crate) async fn handle_txn_chunk(
    ctx: Arc<Context>,
    parts: Parts,
    body: Bytes,
) -> hyper::Result<Response<Body>> {
    let query = parts.uri.query().unwrap_or("");
    let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect();

    let keyspace_id = match get_param::<u32>(&query_pairs, "keyspace_id") {
        None => {
            return Ok(make_response(
                StatusCode::BAD_REQUEST,
                "missing keyspace id",
            ));
        }
        Some(keyspace_id) => keyspace_id,
    };
    let api_version = get_param::<u32>(&query_pairs, "api_version").unwrap_or(TXN_CHUNK_API_V1);
    if !is_valid_txn_chunk_api_version(api_version) {
        return Ok(make_response(
            StatusCode::BAD_REQUEST,
            format!("invalid api_version '{}'", api_version),
        ));
    }

    let Some(_keyspace_permit) = ctx.coprocessor_limiter.acquire_permit(keyspace_id).await else {
        return Ok(make_response(
            StatusCode::SERVICE_UNAVAILABLE,
            "wait permit queue full",
        ));
    };
    let keyspace_info = match ctx
        .txn_chunk_handler
        .acquire_keyspace_info(ctx.clone(), keyspace_id)
        .await
    {
        Ok(keyspace_info) => keyspace_info,
        Err(err) => {
            return Ok(make_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("failed to acquire keyspace info {:?}", err),
            ));
        }
    };

    // TODO: remove this API after client-go is updated.
    if parts.method == http::Method::GET {
        let resp = GetAvailabilityResp { available: true };
        let json = serde_json::to_string(&resp).unwrap();
        return Ok(Response::builder()
            .header(header::CONTENT_TYPE, "application/json")
            .body(json.into())
            .unwrap());
    }

    let chunk_id = match ctx.pd.get_tso().await {
        Ok(tso) => tso.into_inner(),
        Err(err) => {
            return Ok(make_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("failed to create txn chunk file {:?}", err),
            ));
        }
    };
    create_txn_chunk(
        chunk_id,
        ctx.dfs.clone(),
        parts,
        body,
        &keyspace_info,
        ctx.txn_chunk_handler.target_block_size,
        api_version,
    )
    .await
}

#[derive(Default, Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct GetAvailabilityResp {
    pub available: bool,
}

#[derive(Default, Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct CreateTxnChunkResp {
    pub chunk_id: u64,
}

pub(crate) async fn create_txn_chunk(
    chunk_id: u64,
    dfs: Arc<dyn Dfs>,
    parts: Parts,
    body: Bytes,
    keyspace_info: &KeyspaceInfo,
    target_block_size: usize,
    api_version: u32,
) -> hyper::Result<Response<Body>> {
    if parts.method != http::Method::POST {
        return Ok(make_response(StatusCode::BAD_REQUEST, "invalid method"));
    }
    if body.len() < 4 {
        return Ok(make_response(StatusCode::BAD_REQUEST, "body is invalid"));
    }
    let data_len = body.len() - 4;
    let checksum = (&body[data_len..]).get_u32_le();
    let mut body_buf = &body[..data_len];
    if ChecksumType::Crc32.checksum(body_buf) != checksum {
        return Ok(make_response(StatusCode::BAD_REQUEST, "checksum mismatch"));
    }
    let file_format = api_version_to_txn_file_format_version(api_version);
    let mut txn_chunk_builder = TxnChunkBuilder::new(
        chunk_id,
        target_block_size,
        keyspace_info.encryption_key.clone(),
        file_format,
    );
    while !body_buf.is_empty() {
        let key_len = body_buf.get_u16_le() as usize;
        let key = &body_buf[..key_len];
        body_buf.advance(key_len);
        let op: u8 = body_buf.get_u8();
        let pessimistic_action = if api_version >= TXN_CHUNK_API_V2 {
            let action = body_buf.get_u8();
            match PessimisticAction::try_from(action) {
                Ok(a) => a,
                Err(e) => {
                    return Ok(make_response(StatusCode::BAD_REQUEST, e));
                }
            }
        } else {
            // for legacy client that does not contains pessimistic action.
            PessimisticAction::SkipCheck
        };
        let val_len = body_buf.get_u32_le() as usize;
        let val = &body_buf[..val_len];
        body_buf.advance(val_len);
        txn_chunk_builder.add_entry(InnerKey::from_outer_key(key), op, pessimistic_action, val);
    }
    drop(body);
    let mut txn_chunk_buf = vec![];
    txn_chunk_builder.finish(&mut txn_chunk_buf);
    let opts = dfs::Options::default().with_type(dfs::FileType::TxnChunk);
    if let Err(err) = dfs.create(chunk_id, txn_chunk_buf.into(), opts).await {
        return Ok(make_response(
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("failed to create txn chunk file {:?}", err),
        ));
    }
    let resp = CreateTxnChunkResp { chunk_id };
    let json = serde_json::to_string(&resp).unwrap();
    Ok(Response::builder()
        .header(header::CONTENT_TYPE, "application/json")
        .body(json.into())
        .unwrap())
}

#[derive(Clone)]
pub(crate) struct KeyspaceInfo {
    encryption_key: Option<EncryptionKey>,
}

pub(crate) struct TxnChunkHandler {
    target_block_size: usize,
}

impl TxnChunkHandler {
    pub(crate) fn new(target_block_size: usize) -> Self {
        Self { target_block_size }
    }

    async fn acquire_keyspace_info(
        &self,
        ctx: Arc<Context>,
        keyspace_id: u32,
    ) -> Result<KeyspaceInfo> {
        let encryption_key = encryption::load_encryption_key_for_keyspace(
            &ctx.keyspace_encryption_info_cache,
            keyspace_id,
            ctx.pd.clone(),
            ctx.dfs.clone(),
            ctx.encryption_key_manager.clone(),
            &ctx.master_key,
        )
        .await?;

        Ok(KeyspaceInfo { encryption_key })
    }
}

#[cfg(test)]
mod tests {
    use std::{convert::TryInto, ops::Deref, sync::Arc};

    use bytes::{Buf, BufMut};
    use cloud_encryption::EncryptionState;
    use futures::{StreamExt, executor::block_on};
    use http::Method;
    use kvengine::{
        GLOBAL_SHARD_END_KEY, Iterator, UserMeta, dfs,
        dfs::{Dfs, FileType, InMemFs},
        table::{
            InnerKey, TxnChunk, TxnCtx, TxnFile, TxnFileId, TxnFileIterator, file::InMemFile,
            sstable::BlockCache,
        },
    };

    use crate::txn_chunk::{
        CreateTxnChunkResp, KeyspaceInfo, TARGET_BLOCK_SIZE_DEF, create_txn_chunk,
    };

    fn txn_chunk_format_version(chunk_data: &[u8]) -> u16 {
        // Footer layout:
        // u32 index_offset, u32 hash_index_offset, u32 properties_offset, u8 reserved,
        // u8 checksum_type, u16 format_version, u32 magic.
        assert!(chunk_data.len() >= 6);
        u16::from_le_bytes(
            chunk_data[chunk_data.len() - 6..chunk_data.len() - 4]
                .try_into()
                .unwrap(),
        )
    }

    #[test]
    fn test_create_txn_chunk() {
        let mut req_body = vec![];
        let buf = &mut req_body;
        for i in 0..100 {
            let key = format!("key{:03}", i);
            let val = format!("val{:03}", i);
            buf.put_u16_le(key.len() as u16);
            buf.extend_from_slice(key.as_bytes());
            buf.put_u8(1);
            buf.put_u32_le(val.len() as u32);
            buf.extend_from_slice(val.as_bytes());
        }
        let check_sum = crc32fast::hash(buf);
        buf.put_u32_le(check_sum);
        let chunk_id = 155;
        let dfs: Arc<dyn Dfs> = Arc::new(InMemFs::new());
        let req = http::Request::builder()
            .method(Method::POST)
            .body(hyper::Body::from(req_body))
            .unwrap();
        let keyspace_info = KeyspaceInfo {
            encryption_key: None,
        };
        let (parts, body) = req.into_parts();
        let body = block_on(hyper::body::to_bytes(body)).unwrap();
        let mut res = dfs
            .get_runtime()
            .block_on(create_txn_chunk(
                chunk_id,
                dfs.clone(),
                parts,
                body,
                &keyspace_info,
                TARGET_BLOCK_SIZE_DEF,
                super::TXN_CHUNK_API_V1,
            ))
            .unwrap();
        assert!(res.status().is_success());
        let body = dfs
            .get_runtime()
            .block_on(res.body_mut().next())
            .unwrap()
            .unwrap();
        let resp: CreateTxnChunkResp = serde_json::from_slice(body.chunk()).unwrap();
        assert_eq!(resp.chunk_id, 155);
        let opts = dfs::Options::default().with_type(FileType::TxnChunk);
        let chunk_data = dfs
            .get_runtime()
            .block_on(dfs.read_file(chunk_id, opts))
            .unwrap();
        assert!(!chunk_data.is_empty());
        assert_eq!(
            txn_chunk_format_version(&chunk_data),
            kvengine::table::TXN_FILE_FORMAT_V1
        );
        let txn_chunk = TxnChunk::new(
            Arc::new(InMemFile::new(155, chunk_data)),
            BlockCache::None,
            EncryptionState::Disabled,
        )
        .unwrap();
        let user_meta = UserMeta::new(1, 2).to_array().to_vec();
        let lower_bound = InnerKey::from_inner_buf(b"");
        let upper_bound = InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY);
        let txn_ctx = TxnCtx::new(user_meta.into(), vec![].into(), 2, lower_bound, upper_bound);
        let txn_file = TxnFile::new(TxnFileId::new(1, 1, 1), vec![txn_chunk], txn_ctx).unwrap();
        let mut iter = TxnFileIterator::new(txn_file, false);
        iter.rewind();
        let mut i = 0;
        while iter.valid() {
            assert_eq!(iter.key().deref(), format!("key{:03}", i).as_bytes());
            assert_eq!(iter.value().get_value(), format!("val{:03}", i).as_bytes());
            assert_eq!(iter.get_op(), 1);
            iter.next();
            i += 1;
        }
        assert_eq!(i, 100);
    }

    #[test]
    fn test_create_txn_chunk_v1() {
        let mut req_body = vec![];
        let buf = &mut req_body;
        for i in 0..100 {
            let key = format!("key{:03}", i);
            let val = format!("val{:03}", i);
            buf.put_u16_le(key.len() as u16);
            buf.extend_from_slice(key.as_bytes());
            buf.put_u8(1);
            // pessimistic_action
            buf.put_u8(0);
            buf.put_u32_le(val.len() as u32);
            buf.extend_from_slice(val.as_bytes());
        }
        let check_sum = crc32fast::hash(buf);
        buf.put_u32_le(check_sum);
        let chunk_id = 155;
        let dfs: Arc<dyn Dfs> = Arc::new(InMemFs::new());
        let req = http::Request::builder()
            .method(Method::POST)
            .body(hyper::Body::from(req_body))
            .unwrap();
        let keyspace_info = KeyspaceInfo {
            encryption_key: None,
        };
        let (parts, body) = req.into_parts();
        let body = block_on(hyper::body::to_bytes(body)).unwrap();
        let mut res = dfs
            .get_runtime()
            .block_on(create_txn_chunk(
                chunk_id,
                dfs.clone(),
                parts,
                body,
                &keyspace_info,
                TARGET_BLOCK_SIZE_DEF,
                super::TXN_CHUNK_API_V2,
            ))
            .unwrap();
        assert!(res.status().is_success());
        let body = dfs
            .get_runtime()
            .block_on(res.body_mut().next())
            .unwrap()
            .unwrap();
        let resp: CreateTxnChunkResp = serde_json::from_slice(body.chunk()).unwrap();
        assert_eq!(resp.chunk_id, 155);
        let opts = dfs::Options::default().with_type(FileType::TxnChunk);
        let chunk_data = dfs
            .get_runtime()
            .block_on(dfs.read_file(chunk_id, opts))
            .unwrap();
        assert!(!chunk_data.is_empty());
        assert_eq!(
            txn_chunk_format_version(&chunk_data),
            kvengine::table::TXN_FILE_FORMAT_V2
        );
        let txn_chunk = TxnChunk::new(
            Arc::new(InMemFile::new(155, chunk_data)),
            BlockCache::None,
            EncryptionState::Disabled,
        )
        .unwrap();
        let user_meta = UserMeta::new(1, 2).to_array().to_vec();
        let lower_bound = InnerKey::from_inner_buf(b"");
        let upper_bound = InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY);
        let txn_ctx = TxnCtx::new(user_meta.into(), vec![].into(), 2, lower_bound, upper_bound);
        let txn_file = TxnFile::new(TxnFileId::new(1, 1, 1), vec![txn_chunk], txn_ctx).unwrap();
        let mut iter = TxnFileIterator::new(txn_file, false);
        iter.rewind();
        let mut i = 0;
        while iter.valid() {
            assert_eq!(iter.key().deref(), format!("key{:03}", i).as_bytes());
            assert_eq!(iter.value().get_value(), format!("val{:03}", i).as_bytes());
            assert_eq!(iter.get_op(), 1);
            iter.next();
            i += 1;
        }
        assert_eq!(i, 100);
    }
}