use std::{
    collections::HashMap,
    sync::{Arc, atomic::AtomicU64},
};

use log::info;
use mysql_async::{Conn, Transaction, TxOpts, prelude::*};
use tici_common::index::IndexKey;
use tici_grpc::S3Location;

use crate::{
    HandleRange, ShardInfo, ShardInfoUpdate, ShardManifest, ShardStats,
    meta::{metrics::BACKEND_DELAY_DURATION_HISTOGRAM, store::metastore::META_DATABASE_NAME},
};

pub const META_TABLE_NAMES: [&str; 1] = ["tici_shard_meta"];

pub async fn ensure_schema(conn: &mut Conn) -> Result<(), anyhow::Error> {
    info!("Ensure ShardStore shard schema");
    let _s: Vec<String> = conn
        .query(format!(
            r"
            CREATE DATABASE IF NOT EXISTS {0};
            CREATE TABLE IF NOT EXISTS {0}.tici_shard_meta (
                keyspace_id INT UNSIGNED NOT NULL,
                table_id BIGINT UNSIGNED NOT NULL,
                index_id BIGINT UNSIGNED NOT NULL,
                start_key VARBINARY(65535) NOT NULL,
                end_key VARBINARY(65535) NOT NULL,
                shard_id BIGINT UNSIGNED AUTO_INCREMENT,
                epoch BIGINT UNSIGNED NOT NULL,
                seq BIGINT UNSIGNED NOT NULL,
                shard_writer VARCHAR(255) NOT NULL,
                manifest LONGTEXT NOT NULL,
                progress TEXT NOT NULL,
                frag_upload_bucket TEXT NOT NULL,
                frag_upload_prefix TEXT NOT NULL,
                commit_leader_term BIGINT UNSIGNED NOT NULL,
                PRIMARY KEY (shard_id),
                KEY idx_start_key (keyspace_id,table_id,index_id,start_key(1024))
            );",
            META_DATABASE_NAME
        ))
        .await?;
    info!("Ensure ShardStore shard schema finished");
    Ok(())
}

#[allow(clippy::type_complexity)]
pub async fn restore(conn: &mut Conn) -> Result<Vec<ShardInfo>, anyhow::Error> {
    #[cfg(not(any(test, feature = "testexport")))]
    info!("restoring shards from meta table");
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let sql = format!(
        r"SELECT keyspace_id,table_id,index_id,shard_id,start_key,end_key,epoch,seq,manifest,shard_writer,progress,frag_upload_bucket,frag_upload_prefix FROM {}.tici_shard_meta",
        META_DATABASE_NAME
    );
    let shards: Vec<ShardInfo> = tx
        .exec_iter(sql, ())
        .await?
        .map_and_drop(|row: mysql_async::Row| {
            // Extract fields by index
            let keyspace_id: u32 = row.get("keyspace_id").unwrap_or_default();
            let table_id: i64 = row.get("table_id").unwrap_or_default();
            let index_id: i64 = row.get("index_id").unwrap_or_default();
            let shard_id: u64 = row.get("shard_id").unwrap_or_default();
            let start_key: Vec<u8> = row.get("start_key").unwrap_or_default();
            let end_key: Vec<u8> = row.get("end_key").unwrap_or_default();
            let epoch: u64 = row.get("epoch").unwrap_or_default();
            let seq: u64 = row.get("seq").unwrap_or_default();
            let manifest: String = row.get("manifest").unwrap_or_default();
            let shard_writer: String = row.get("shard_writer").unwrap_or_default();
            let progress: String = row.get("progress").unwrap_or_default();
            let frag_upload_bucket: String = row.get("frag_upload_bucket").unwrap_or_default();
            let frag_upload_prefix: String = row.get("frag_upload_prefix").unwrap_or_default();

            let mut proc: crate::meta::writer::ShardBuildProgress =
                serde_json::from_str(progress.as_str()).unwrap_or_default();
            proc.mark_cdc_last_file_persisted();
            let manifest = match serde_json::from_str(manifest.as_str()) {
                Ok(e) => e,
                Err(_) => {
                    // Only happens when a new index is bootstrapped,
                    // and the meta service restored before the shard writer has reported its
                    // first heartbeat.
                    ShardManifest::new_for_first()
                }
            };
            let stats = ShardStats::new(
                manifest.get_total_estimated_bytes(),
                manifest.get_fragments().len() as u64,
            );
            ShardInfo {
                index_key: IndexKey::new(keyspace_id, table_id, index_id),
                shard_id,
                range: HandleRange::new(start_key, end_key),
                epoch,
                seq,
                manifest,
                writer: shard_writer,
                progress: proc,
                frag_upload_location: S3Location {
                    bucket: frag_upload_bucket,
                    prefix: frag_upload_prefix,
                },
                stats,
                append_empty_persist_count: Arc::new(AtomicU64::new(0)),
            }
        })
        .await?;

    info!("restored {} shards from meta table", shards.len());
    Ok(shards)
}

pub async fn init(
    conn: &mut Conn,
    new_shard: &ShardInfo,
    leader_term: i64,
) -> Result<u64, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    let shard_id = match inner_add_shard("init_shard", &mut tx, new_shard, leader_term).await {
        Ok(x) => x,
        Err(e) => {
            tx.rollback().await?;
            return Err(e);
        }
    };

    tx.commit().await?;
    Ok(shard_id)
}

pub async fn split_shards(
    conn: &mut Conn,
    original_shard: &ShardInfo,
    new_shard_infos: Vec<ShardInfo>,
    leader_term: i64,
    batch: usize,
) -> Result<Vec<ShardInfo>, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    match inner_remove_shard("split_shards_remove", &mut tx, true, original_shard).await {
        Ok(_) => (),
        Err(e) => {
            tx.rollback().await?;
            return Err(e);
        }
    };

    let mut new_shard_infos = new_shard_infos;
    let chunk_size = if batch == 0 { 1000 } else { batch };
    for (idx, shards) in new_shard_infos.chunks_mut(chunk_size).enumerate() {
        if chunk_size > 100 && idx > 0 {
            log::info!("batch split shards: handle {}", idx);
        }
        let shard_refs: Vec<&ShardInfo> = shards.iter().collect();
        let shard_ids =
            match inner_add_shards("split_shards_add", &mut tx, &shard_refs, leader_term).await {
                Ok(ids) => ids,
                Err(e) => {
                    tx.rollback().await?;
                    return Err(e);
                }
            };
        for (shard_info, shard_id) in shards.iter_mut().zip(shard_ids.into_iter()) {
            shard_info.shard_id = shard_id;
        }
    }

    tx.commit().await?;
    Ok(new_shard_infos)
}

pub async fn update_shard(
    conn: &mut Conn,
    tag: &str,
    shard: &ShardInfo,
    shard_update: ShardInfoUpdate,
) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let mut changed = false;
    let progress = shard_update.progress;
    let manifest = shard_update.manifest;
    let new_seq = shard_update.seq;
    let writer = shard_update.writer;
    let mock_delayed: u64 = (|| {
        fail::fail_point!("tidb_backend_mock_delayed_update", |t| {
            let Some(t) = t else {
                return 0;
            };
            match t.parse::<u64>() {
                Ok(v) => v,
                Err(e) => {
                    log::warn!(
                        "failpoint tidb_backend_mock_delayed_update: invalid param {t:?}: {e}"
                    );
                    0
                }
            }
        });
        0
    })();
    if mock_delayed > 0 {
        tokio::time::sleep(std::time::Duration::from_millis(mock_delayed)).await;
    }

    let (progress_updater, progress_json) = if let Some(p) = &progress {
        changed |= p != &shard.progress;
        ("progress=:progress".to_string(), serde_json::to_string(p)?)
    } else {
        (String::new(), String::new())
    };
    let (manifest_updater, manifest_json) = if let Some(m) = &manifest {
        changed |= m != &shard.manifest;
        ("manifest=:manifest".to_string(), serde_json::to_string(m)?)
    } else {
        (String::new(), String::new())
    };
    let (seq_updater, new_seq) = if let Some(s) = &new_seq {
        changed |= s != &shard.seq;
        ("seq=:seq".to_string(), *s)
    } else {
        (String::new(), 0)
    };
    let (writer_updater, new_writer) = if let Some(s) = &writer {
        changed |= s != &shard.writer;
        ("shard_writer=:writer".to_string(), s.clone())
    } else {
        (String::new(), String::new())
    };

    if !changed {
        // No change has been made, no need to update.
        return Ok(());
    }
    let updater = [
        progress_updater,
        manifest_updater,
        seq_updater,
        writer_updater,
    ]
    .into_iter()
    .filter(|x| !x.is_empty())
    .collect::<Vec<_>>()
    .join(",");
    let sql = format!(
        r"UPDATE {}.tici_shard_meta SET {} WHERE keyspace_id=:keyspace_id and table_id=:table_id and index_id=:index_id and start_key=:where_start_key and epoch=:where_epoch and seq=:where_seq and shard_id=:where_id;",
        META_DATABASE_NAME, updater
    );
    let q = tx
        .exec_iter(
            sql,
            params! {
                "keyspace_id" => shard.keyspace_id(),
                "table_id" => shard.table_id(),
                "index_id" => shard.index_id(),
                "seq" => new_seq,
                "progress" => progress_json,
                "manifest" => manifest_json,
                "writer" => &new_writer,
                "where_start_key" => shard.range.encoded_start(),
                "where_epoch" => shard.epoch,
                "where_seq" => shard.seq,
                "where_id" => shard.shard_id,
            },
        )
        .await?;

    if q.affected_rows() != 1 {
        let affect_rows = q.affected_rows();
        tx.rollback().await?;
        return Err(anyhow::format_err!(
            "update_shard({:?}): detected inconsistency, shard={}, range={:?}, epoch={:?}, seq={:?}, progress={:?}, manifest={:?}, writer={:?}, affected_rows={}",
            tag,
            shard.to_long_ident(),
            shard.range,
            shard.epoch,
            shard.seq,
            progress,
            manifest,
            new_writer,
            affect_rows,
        ));
    }
    tx.commit().await?;
    Ok(())
}

pub async fn merge_shards(
    conn: &mut Conn,
    original_shards: Vec<&ShardInfo>,
    new_shard: ShardInfo,
    leader_term: i64,
) -> Result<ShardInfo, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    for shard in original_shards.iter() {
        match inner_remove_shard("merge_shards_remove", &mut tx, true, shard).await {
            Ok(_) => (),
            Err(e) => {
                tx.rollback().await?;
                return Err(e);
            }
        };
    }

    let mut new_shard = new_shard;
    let shard_id = match inner_add_shard("merge_shards_add", &mut tx, &new_shard, leader_term).await
    {
        Ok(x) => x,
        Err(e) => {
            tx.rollback().await?;
            return Err(e);
        }
    };
    new_shard.shard_id = shard_id;

    tx.commit().await?;
    Ok(new_shard)
}

pub async fn remove_shard(conn: &mut Conn, shard: &ShardInfo) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    inner_remove_shard("remove_shard", &mut tx, false, shard).await?;
    tx.commit().await?;
    Ok(())
}

async fn inner_remove_shard(
    tag: &str,
    tx: &mut Transaction<'_>,
    must_exist: bool,
    shard: &ShardInfo,
) -> Result<(), anyhow::Error> {
    let delete_sql = tx.exec_iter(
        format!(r"DELETE FROM {}.tici_shard_meta WHERE shard_id=:shard_id and start_key=:start_key and epoch=:epoch and seq=:seq;", META_DATABASE_NAME),
                params! {
                    "shard_id" => shard.shard_id,
                    "start_key" => shard.range.encoded_start(),
                    "epoch" => shard.epoch,
                    "seq" => shard.seq,
                }
            ).await?;

    if delete_sql.affected_rows() != 1 && must_exist {
        let affect_rows = delete_sql.affected_rows();
        return Err(anyhow::format_err!(
            "{}: detected inconsistency, shard={}, affected_rows={}",
            tag,
            shard.to_long_ident(),
            affect_rows,
        ));
    }
    Ok(())
}

async fn inner_add_shard(
    tag: &str,
    tx: &mut Transaction<'_>,
    new_shard: &ShardInfo,
    leader_term: i64,
) -> Result<u64, anyhow::Error> {
    let v = vec![new_shard];
    let res = inner_add_shards(tag, tx, v.as_slice(), leader_term).await?;
    Ok(res[0])
}

async fn inner_add_shards(
    tag: &str,
    tx: &mut Transaction<'_>,
    new_shards: &[&ShardInfo],
    leader_term: i64,
) -> Result<Vec<u64>, anyhow::Error> {
    // NOTE `shard_id` is not guaranteed to be increment, although it is marked as
    // AUTO_IMCREMENT in the storage table. We have to use `epoch` to get the
    // partial ordering of two shards if they overlapped with each other.
    if new_shards.is_empty() {
        return Ok(vec![]);
    }

    let mut sql = format!(
        "INSERT IGNORE INTO {}.tici_shard_meta \
        (keyspace_id,table_id,index_id,start_key,end_key,epoch,seq,shard_writer,manifest,commit_leader_term,frag_upload_bucket,frag_upload_prefix,progress) VALUES ",
        META_DATABASE_NAME
    );

    let mut values_sql = Vec::new();

    let mut params_map: HashMap<Vec<u8>, mysql_async::Value> = HashMap::new();

    for (i, shard) in new_shards.iter().enumerate() {
        let manifest_json = serde_json::to_string(&shard.manifest)?;
        let progress_json = serde_json::to_string(&shard.progress)?;

        values_sql.push(format!(
            "(:keyspace_id_{i},:table_id_{i},:index_id_{i},:start_key_{i},:end_key_{i},:epoch_{i},:seq_{i},:shard_writer_{i},:manifest_{i},:commit_leader_term_{i},:frag_upload_bucket_{i},:frag_upload_prefix_{i},:progress_{i})"
        ));

        params_map.insert(
            format!("keyspace_id_{i}").into_bytes(),
            shard.keyspace_id().into(),
        );
        params_map.insert(
            format!("table_id_{i}").into_bytes(),
            shard.table_id().into(),
        );
        params_map.insert(
            format!("index_id_{i}").into_bytes(),
            shard.index_id().into(),
        );
        params_map.insert(
            format!("start_key_{i}").into_bytes(),
            shard.range.encoded_start().into(),
        );
        params_map.insert(
            format!("end_key_{i}").into_bytes(),
            shard.range.encoded_end().into(),
        );
        params_map.insert(format!("epoch_{i}").into_bytes(), shard.epoch.into());
        params_map.insert(format!("seq_{i}").into_bytes(), shard.seq.into());
        params_map.insert(
            format!("shard_writer_{i}").into_bytes(),
            shard.writer.clone().into(),
        );
        params_map.insert(
            format!("manifest_{i}").into_bytes(),
            manifest_json.clone().into(),
        );
        params_map.insert(
            format!("commit_leader_term_{i}").into_bytes(),
            leader_term.into(),
        );
        params_map.insert(
            format!("frag_upload_bucket_{i}").into_bytes(),
            shard.frag_upload_location.bucket.clone().into(),
        );
        params_map.insert(
            format!("frag_upload_prefix_{i}").into_bytes(),
            shard.frag_upload_location.prefix.clone().into(),
        );
        params_map.insert(
            format!("progress_{i}").into_bytes(),
            progress_json.clone().into(),
        );
    }

    sql.push_str(&values_sql.join(","));

    let params = mysql_async::Params::Named(params_map);

    let q = {
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["inner_add_shards_insert"])
            .start_timer();
        tx.exec_iter(sql, params).await?
    };

    if q.affected_rows() != new_shards.len() as u64 {
        let affect_rows = q.affected_rows();
        return Err(anyhow::format_err!(
            "add_shards({}): detected inconsistency, to_add_shards={:?} affected_rows={}",
            tag,
            new_shards,
            affect_rows,
        ));
    }

    let mut shard_ids = Vec::with_capacity(new_shards.len());
    let mut sql = format!(
        "SELECT keyspace_id, table_id, index_id, start_key, shard_id FROM {}.tici_shard_meta WHERE ",
        META_DATABASE_NAME
    );
    let mut conds = Vec::with_capacity(new_shards.len());
    let mut params_map: HashMap<Vec<u8>, mysql_async::Value> = HashMap::new();
    for (i, shard) in new_shards.iter().enumerate() {
        conds.push(format!(
            "(keyspace_id=:keyspace_id_{i} AND table_id=:table_id_{i} AND index_id=:index_id_{i} AND start_key=:start_key_{i})"
        ));
        params_map.insert(
            format!("keyspace_id_{i}").into_bytes(),
            shard.keyspace_id().into(),
        );
        params_map.insert(
            format!("table_id_{i}").into_bytes(),
            shard.table_id().into(),
        );
        params_map.insert(
            format!("index_id_{i}").into_bytes(),
            shard.index_id().into(),
        );
        params_map.insert(
            format!("start_key_{i}").into_bytes(),
            shard.range.encoded_start().into(),
        );
    }
    let params = mysql_async::Params::Named(params_map);
    sql.push_str(&conds.join(" OR "));

    let rows: Vec<(u32, i64, i64, Vec<u8>, u64)> = {
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["inner_add_shards_select"])
            .start_timer();
        tx.exec_iter(sql, params)
            .await?
            .map_and_drop(mysql_async::from_row::<(u32, i64, i64, Vec<u8>, u64)>)
            .await?
    };

    // This happens when there are duplicated `start_key`s.
    if rows.len() != new_shards.len() {
        return Err(anyhow::format_err!(
            "add_shards({}): detected inconsistency, to_add_shard: [{}]{:?} rows: [{}]{:?}",
            tag,
            new_shards.len(),
            new_shards,
            rows.len(),
            rows,
        ));
    }

    let mut id_map = HashMap::with_capacity(rows.len());
    for (keyspace_id, table_id, index_id, start_key, shard_id) in rows {
        id_map.insert((keyspace_id, table_id, index_id, start_key), shard_id);
    }

    if id_map.len() != new_shards.len() {
        return Err(anyhow::format_err!(
            "add_shards({}): detected inconsistency, to_add_shard: [{}]{:?} id_map: [{}]{:?}",
            tag,
            new_shards.len(),
            new_shards,
            id_map.len(),
            id_map,
        ));
    }

    for shard in new_shards {
        let key = (
            shard.keyspace_id(),
            shard.table_id(),
            shard.index_id(),
            shard.range.encoded_start(),
        );
        match id_map.get(&key) {
            Some(&shard_id) => {
                shard_ids.push(shard_id);
            }
            None => {
                return Err(anyhow::format_err!(
                    "add_shards({}): detected inconsistency, to_add_shard={:?} can't find={:?}",
                    tag,
                    shard,
                    key,
                ));
            }
        }
    }

    Ok(shard_ids)
}