use std::{
    collections::HashSet,
    sync::{
        Mutex,
        atomic::{AtomicU64, Ordering},
    },
};

use async_trait::async_trait;
use tici_common::index::IndexKey;
use tici_grpc::{IndexInfo, TableInfo};
use tici_import::{
    ImportJobID, ImportJobMeta, ImportTask, ImportTaskID, ImportTaskParams, ImportTaskSummary,
    TiDBImportTaskID,
};
use tici_service_discovery::election::leadership::{Leadership, MockLeadership};

use super::{
    IndexMetaProperties, MetaStore,
    import_task::{ImportJobStat, InactiveImportTask},
    metastore::ImportIndexProgress,
};
use crate::{ShardInfo, ShardInfoUpdate};

#[allow(dead_code)]
#[derive(Default)]
pub(crate) struct DummyMetaStore {
    leadership: MockLeadership,
}

#[allow(dead_code)]
#[derive(Default)]
pub(crate) struct RecordingMetaStore {
    leadership: MockLeadership,
    pub(crate) next_shard_id: AtomicU64,
    pub(crate) removed_shard_ids: Mutex<Vec<u64>>,
    pub(crate) shard_updates: Mutex<Vec<(String, ShardInfoUpdate)>>,
    pub(crate) next_import_tasks: Mutex<Vec<ImportTask>>,
    pub(crate) running_import_tasks: Mutex<Vec<ImportTask>>,
    pub(crate) reclaimed_import_tasks: Mutex<Vec<(ImportTaskID, String)>>,
    pub(crate) reclaimed_import_workers: Mutex<Vec<String>>,
}

#[async_trait]
impl MetaStore for DummyMetaStore {
    async fn restore_indexes(
        &self,
    ) -> Result<Vec<(IndexInfo, TableInfo, String, String, IndexMetaProperties)>, anyhow::Error>
    {
        unimplemented!()
    }

    async fn init_index(
        &self,
        _table_info: &TableInfo,
        _index_info: &IndexInfo,
        _cdc_s3_bucket: String,
        _cdc_s3_prefix: String,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn finish_index_append_base(&self, _index_key: &IndexKey) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn drop_index(&self, _index_key: &IndexKey) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn restore_shards(&self) -> Result<Vec<ShardInfo>, anyhow::Error> {
        Ok(vec![])
    }

    async fn init_shard(&self, _new_shard: &ShardInfo) -> Result<u64, anyhow::Error> {
        unimplemented!()
    }

    async fn remove_shard(&self, _shard: &ShardInfo) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn split_shards(
        &self,
        _original_shard: &ShardInfo,
        _new_shard_infos: Vec<ShardInfo>,
        _batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        unimplemented!()
    }

    async fn merge_shards(
        &self,
        _original_shards: Vec<&ShardInfo>,
        _new_shard: ShardInfo,
    ) -> Result<ShardInfo, anyhow::Error> {
        unimplemented!()
    }

    async fn update_shard(
        &self,
        _tag: &str,
        _shard: &ShardInfo,
        _shard_update: ShardInfoUpdate,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn create_or_get_import_job(
        &self,
        _job_meta: &ImportJobMeta,
    ) -> Result<ImportJobID, anyhow::Error> {
        unimplemented!()
    }

    async fn finish_import_job_upload(
        &self,
        _tidb_task_id: &TiDBImportTaskID,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn finish_import_job<'a>(
        &self,
        _job_ids: &'a [ImportJobID],
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn fail_import_job<'a>(&self, _job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn get_jobs_should_finish(&self) -> Result<Vec<ImportJobStat>, mysql_async::Error> {
        unimplemented!()
    }

    async fn get_running_import_task_count(&self) -> Result<usize, mysql_async::Error> {
        Ok(0)
    }

    async fn get_running_import_tasks(&self) -> Result<Vec<ImportTask>, mysql_async::Error> {
        Ok(vec![])
    }

    async fn get_import_index_progress(
        &self,
        _keyspace_id: u32,
        _table_id: i64,
        _index_id: i64,
    ) -> Result<Option<ImportIndexProgress>, anyhow::Error> {
        Ok(None)
    }

    async fn get_active_import_jobs(&self) -> Result<Vec<ImportJobMeta>, anyhow::Error> {
        Ok(vec![])
    }

    async fn get_external_terminal_task_keys(
        &self,
        _task_keys: &[String],
    ) -> Result<HashSet<String>, anyhow::Error> {
        Ok(HashSet::new())
    }

    async fn add_import_task(
        &self,
        _tidb_task_id: &TiDBImportTaskID,
        _params: ImportTaskParams,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn next_import_task(
        &self,
        _owner_addr: &str,
    ) -> Result<Option<ImportTask>, mysql_async::Error> {
        unimplemented!()
    }

    async fn reclaim_import_task(
        &self,
        _task_id: ImportTaskID,
        _owner_addr: &str,
    ) -> Result<bool, anyhow::Error> {
        unimplemented!()
    }

    async fn reclaim_import_tasks_by_worker(
        &self,
        _owner_addr: &str,
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        Ok(vec![])
    }

    async fn fail_import_jobs_and_reclaim_running_tasks(
        &self,
        _job_ids: &[ImportJobID],
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        Ok(vec![])
    }

    async fn inactive_import_tasks(
        &self,
        _inactive_secs: u64,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        unimplemented!()
    }

    async fn reset_inactive_import_tasks(
        &self,
        _inactive_secs: u64,
        _limit: usize,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        unimplemented!()
    }

    async fn import_task_build_success(
        &self,
        _task_id: ImportTaskID,
        _owner_addr: &str,
        _owner_seq: u64,
        _summary: ImportTaskSummary,
    ) -> Result<bool, anyhow::Error> {
        unimplemented!()
    }

    async fn import_task_build_fail(
        &self,
        _task_id: ImportTaskID,
        _owner_addr: &str,
        _owner_seq: u64,
        _err_msg: &str,
    ) -> Result<bool, anyhow::Error> {
        unimplemented!()
    }

    fn get_leadership(&self) -> &dyn Leadership {
        &self.leadership
    }
}

#[async_trait]
impl MetaStore for RecordingMetaStore {
    async fn restore_indexes(
        &self,
    ) -> Result<Vec<(IndexInfo, TableInfo, String, String, IndexMetaProperties)>, anyhow::Error>
    {
        unimplemented!()
    }

    async fn init_index(
        &self,
        _table_info: &TableInfo,
        _index_info: &IndexInfo,
        _cdc_s3_bucket: String,
        _cdc_s3_prefix: String,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn finish_index_append_base(&self, _index_key: &IndexKey) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn drop_index(&self, _index_key: &IndexKey) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn restore_shards(&self) -> Result<Vec<ShardInfo>, anyhow::Error> {
        Ok(vec![])
    }

    async fn init_shard(&self, _new_shard: &ShardInfo) -> Result<u64, anyhow::Error> {
        Ok(self.next_shard_id.fetch_add(1, Ordering::Relaxed) + 1)
    }

    async fn remove_shard(&self, shard: &ShardInfo) -> Result<(), anyhow::Error> {
        let mut removed_shard_ids = match self.removed_shard_ids.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        removed_shard_ids.push(shard.shard_id);
        Ok(())
    }

    async fn split_shards(
        &self,
        _original_shard: &ShardInfo,
        new_shard_infos: Vec<ShardInfo>,
        _batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        let mut out = new_shard_infos;
        for shard in &mut out {
            shard.shard_id = self.next_shard_id.fetch_add(1, Ordering::Relaxed) + 1;
        }
        Ok(out)
    }

    async fn merge_shards(
        &self,
        _original_shards: Vec<&ShardInfo>,
        mut new_shard: ShardInfo,
    ) -> Result<ShardInfo, anyhow::Error> {
        new_shard.shard_id = self.next_shard_id.fetch_add(1, Ordering::Relaxed) + 1;
        Ok(new_shard)
    }

    async fn update_shard(
        &self,
        tag: &str,
        _shard: &ShardInfo,
        shard_update: ShardInfoUpdate,
    ) -> Result<(), anyhow::Error> {
        let mut shard_updates = match self.shard_updates.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        shard_updates.push((tag.to_string(), shard_update));
        Ok(())
    }

    async fn create_or_get_import_job(
        &self,
        _job_meta: &ImportJobMeta,
    ) -> Result<ImportJobID, anyhow::Error> {
        unimplemented!()
    }

    async fn finish_import_job_upload(
        &self,
        _tidb_task_id: &TiDBImportTaskID,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn finish_import_job<'a>(
        &self,
        _job_ids: &'a [ImportJobID],
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn fail_import_job<'a>(&self, _job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn get_jobs_should_finish(&self) -> Result<Vec<ImportJobStat>, mysql_async::Error> {
        unimplemented!()
    }

    async fn get_running_import_task_count(&self) -> Result<usize, mysql_async::Error> {
        let tasks = match self.running_import_tasks.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        Ok(tasks.len())
    }

    async fn get_running_import_tasks(&self) -> Result<Vec<ImportTask>, mysql_async::Error> {
        let tasks = match self.running_import_tasks.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        Ok(tasks)
    }

    async fn get_import_index_progress(
        &self,
        _keyspace_id: u32,
        _table_id: i64,
        _index_id: i64,
    ) -> Result<Option<ImportIndexProgress>, anyhow::Error> {
        Ok(None)
    }

    async fn get_active_import_jobs(&self) -> Result<Vec<ImportJobMeta>, anyhow::Error> {
        Ok(vec![])
    }

    async fn get_external_terminal_task_keys(
        &self,
        _task_keys: &[String],
    ) -> Result<HashSet<String>, anyhow::Error> {
        Ok(HashSet::new())
    }

    async fn add_import_task(
        &self,
        _tidb_task_id: &TiDBImportTaskID,
        _params: ImportTaskParams,
    ) -> Result<(), anyhow::Error> {
        unimplemented!()
    }

    async fn next_import_task(
        &self,
        owner_addr: &str,
    ) -> Result<Option<ImportTask>, mysql_async::Error> {
        let mut queued = match self.next_import_tasks.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        if queued.is_empty() {
            return Ok(None);
        }
        let mut task = queued.remove(0);
        task.owner_addr = owner_addr.to_string();
        drop(queued);

        let mut running = match self.running_import_tasks.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        running.push(task.clone());
        Ok(Some(task))
    }

    async fn reclaim_import_task(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
    ) -> Result<bool, anyhow::Error> {
        let mut reclaimed = match self.reclaimed_import_tasks.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        reclaimed.push((task_id, owner_addr.to_string()));
        Ok(true)
    }

    async fn reclaim_import_tasks_by_worker(
        &self,
        owner_addr: &str,
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        let task_ids = match self.running_import_tasks.lock() {
            Ok(guard) => guard
                .iter()
                .filter(|task| task.owner_addr == owner_addr)
                .map(|task| task.task_id)
                .collect::<Vec<_>>(),
            Err(poisoned) => poisoned
                .into_inner()
                .iter()
                .filter(|task| task.owner_addr == owner_addr)
                .map(|task| task.task_id)
                .collect::<Vec<_>>(),
        };
        let mut reclaimed = match self.reclaimed_import_workers.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        reclaimed.push(owner_addr.to_string());
        Ok(task_ids)
    }

    async fn fail_import_jobs_and_reclaim_running_tasks(
        &self,
        _job_ids: &[ImportJobID],
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        Ok(vec![])
    }

    async fn inactive_import_tasks(
        &self,
        _inactive_secs: u64,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        unimplemented!()
    }

    async fn reset_inactive_import_tasks(
        &self,
        _inactive_secs: u64,
        _limit: usize,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        unimplemented!()
    }

    async fn import_task_build_success(
        &self,
        _task_id: ImportTaskID,
        _owner_addr: &str,
        _owner_seq: u64,
        _summary: ImportTaskSummary,
    ) -> Result<bool, anyhow::Error> {
        Ok(true)
    }

    async fn import_task_build_fail(
        &self,
        _task_id: ImportTaskID,
        _owner_addr: &str,
        _owner_seq: u64,
        _err_msg: &str,
    ) -> Result<bool, anyhow::Error> {
        Ok(true)
    }

    fn get_leadership(&self) -> &dyn Leadership {
        &self.leadership
    }
}