use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
    time::SystemTime,
};

use either::Either;
use log::{info, warn};
use tici::{HandleRange, RangeBound, ShardInfo, worker::mocks::make_frag_meta_extended};
use tici_common::{
    codec::{previous_key_helper, table_end_key},
    index::IndexKey,
};
use tici_system::readable::ReadableSize;
use tici_test_utils::test_utils::{test_suite::E2ETestSuite, *};
use tokio::sync::RwLock;

mod long_term;

#[derive(Debug, Clone)]
pub struct Workload {
    pub index_key: IndexKey,
    pub default_size_batch: u64,
    pub default_rows_batch: u64,
    pub n_default: u64,
    pub n_parallel: u64,
    pub n_presplit: u64,
    pub n_runner: u64,
    pub update_delay_ms: u64,
    pub do_check: bool,
}

impl Workload {
    pub fn new(index_key: IndexKey, default_size_batch: u64, default_rows_batch: u64) -> Self {
        let n_default = std::env::var("TEST_BENCH_ROWS")
            .map(|v| v.parse().unwrap())
            .unwrap_or(400u64);
        let n_parallel = std::env::var("TEST_PARALLEL")
            .map(|v| v.parse().unwrap())
            .unwrap_or(2u64);
        let n_presplit = std::env::var("TEST_PRESPLIT")
            .map(|v| v.parse().unwrap())
            .unwrap_or(1u64);
        let do_check = std::env::var("TEST_CHECK")
            .map(|v| v.parse().unwrap())
            .unwrap_or(true);
        let n_runner = std::env::var("TEST_RUNNER")
            .map(|v| v.parse().unwrap())
            .unwrap_or(4u64);
        let update_delay_ms = std::env::var("TEST_UPDATE_DELAY")
            .map(|v| v.parse().unwrap())
            .unwrap_or(0u64);
        Self {
            index_key,
            default_size_batch,
            default_rows_batch,
            n_default,
            n_parallel,
            n_presplit,
            n_runner,
            update_delay_ms,
            do_check,
        }
    }

    #[allow(dead_code)] // not used under not linux
    pub fn to_debug_string(&self) -> String {
        format!(
            "(bench_rows={}, parallel={} pre-split={} runner={})",
            self.n_default, self.n_parallel, self.n_presplit, self.n_runner
        )
    }

    fn build_skewed_frag_start(&self, op: u64) -> Vec<u8> {
        // Build a monotonic tail key near table end so tail shards become hotspots.
        let tail_prefix = previous_key_helper(table_end_key(&self.index_key).as_slice(), 1);
        let seq_base = u64::MAX - self.n_default.saturating_sub(1);
        let seq = seq_base.saturating_add(op);

        let mut frag_start = tail_prefix;
        frag_start.extend_from_slice(&seq.to_be_bytes());
        frag_start.push(0);
        frag_start
    }

    fn build_skewed_frag_range(&self, shard: &ShardInfo, frag_start: Vec<u8>) -> HandleRange {
        let mut frag_end = frag_start[..frag_start.len() - 1].to_vec();
        frag_end.push(255);

        if !shard.range.contains_key(&RangeBound::Key(frag_end.clone())) {
            let shard_end = shard.range.encoded_end();
            let mut extra = frag_start.len().saturating_sub(shard_end.len());
            let mut attempts = 0;
            loop {
                let candidate = previous_key_helper(shard_end.as_slice(), extra);
                if candidate > frag_start {
                    frag_end = candidate;
                    break;
                }
                extra = extra.saturating_add(1);
                attempts += 1;
                if attempts > 16 {
                    panic!(
                        "skewed append can't pick frag_end, start={:?}, shard_end={:?}",
                        frag_start, shard_end
                    );
                }
            }
        }
        HandleRange::new(frag_start, frag_end)
    }

    async fn append(&self, suite: Arc<RwLock<E2ETestSuite>>, unit: u64, i: u64) {
        info!("append start {}/{}", unit, i);
        let append_addr = suite.read().await.get_one_worker_node_addr();
        // fragment 0 of index i
        let suite = suite.read().await;
        let append_addr_cloned = append_addr.clone();
        let prefix = format!("a_{}", unit);
        let mock = suite.get_mock_worker(&append_addr_cloned).await;
        let (_, frag_metas) = mock
            .mock_new_frags(
                Either::Left(prefix.clone()),
                self.default_size_batch,
                self.default_rows_batch,
                4,
            )
            .await;
        let example_frag_meta = frag_metas[0].frag_path.clone();
        let default_size_batch = self.default_size_batch;
        let default_rows_batch = self.default_rows_batch;
        let start = SystemTime::now();
        let r = if self.do_check {
            suite
                .must_append_delta_by_mocker(
                    mock,
                    Either::Right(example_frag_meta),
                    15,
                    AppendExtraParam {
                        do_manifest_check: false,
                        upload_to_s3: false,
                        ..Default::default()
                    },
                    default_size_batch,
                    default_rows_batch,
                )
                .await
        } else {
            suite
                .must_append_delta_by_mocker(
                    mock,
                    Either::Right(example_frag_meta),
                    15,
                    AppendExtraParam {
                        do_manifest_check: false,
                        upload_to_s3: false,
                        call_check_append_frag: false,
                        ..Default::default()
                    },
                    default_size_batch,
                    default_rows_batch,
                )
                .await
        };
        let suppose_shard_id = match &r {
            Ok(real_res) => real_res.shard_id,
            Err(e) => e.shard_id(),
        };
        if let Err(e) = r {
            panic!(
                "unexpected add shard error={:?}, suppose_shard_id={}",
                e, suppose_shard_id
            );
        }
        info!(
            "append end {}/{} @ {}, elapsed={}ms",
            unit,
            i,
            suppose_shard_id,
            SystemTime::now().duration_since(start).unwrap().as_millis()
        );
        if unit == 0 && i % 100 == 0 {
            let size = suite.compute_total_size(&self.index_key).await;
            let rows = suite.compute_total_rows(&self.index_key).await;
            info!(
                "current[{}/{}]: total_size={:.6} MiB total_count={} index_stats: ({}) elapsed={:?} s",
                i,
                self.n_default,
                ReadableSize(size).as_mb_f64(),
                rows,
                suite
                    .get_index(&self.index_key)
                    .await
                    .stats
                    .to_short_string(),
                suite.start.elapsed().unwrap().as_secs(),
            );
        }
        info!("append finish {}", i);
    }

    /// A skewed "tail append" workload:
    /// - append to a small range near the overall right edge, so tail shards
    ///   become hotspots. The hotspot shard will move when split/merge happens.
    async fn append_skewed(
        &self,
        suite: Arc<RwLock<E2ETestSuite>>,
        unit: u64,
        op: u64,
        retry_stats: Arc<Mutex<HashMap<i32, u64>>>,
    ) -> u64 {
        let mut att = 0;
        loop {
            info!("skewed append start {}/{}", unit, op);
            let suite = suite.read().await;

            let index = suite.get_index(&self.index_key).await;

            let frag_start = self.build_skewed_frag_start(op);
            let start_bound = RangeBound::Key(frag_start.clone());

            // Find the shard that contains the start key (tail).
            let shard = match index.find_shard_by_key(&start_bound).await {
                Some(s) => s.read().await.clone(),
                None => {
                    att += 1;
                    if att > 10 {
                        panic!(
                            "skewed append can't find shard for tail key, index={:?}",
                            self.index_key
                        );
                    }
                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                    continue;
                }
            };

            let frag_range = self.build_skewed_frag_range(&shard, frag_start);
            let frag_path = mock_frag_path(format!("skew_a_{}_{}", unit, op));
            let frag_meta = make_frag_meta_extended(
                frag_path.as_str(),
                &frag_range,
                self.default_size_batch,
                self.default_rows_batch,
                4,
            );

            let start = SystemTime::now();
            let (append_req, resp) = match suite
                .do_append_delta_frag(
                    shard.writer.as_str(),
                    shard.shard_id,
                    &shard.index_key,
                    create_cdc_file_name(15),
                    vec![frag_meta],
                )
                .await
            {
                Ok(r) => r,
                Err(e) => {
                    att += 1;
                    if att == 1 {
                        warn!(
                            "skewed append retryable rpc error: err={e:?}, shard_id={}, writer={}, range={:?}",
                            shard.shard_id, shard.writer, shard.range,
                        );
                    }
                    if att > 10 {
                        panic!(
                            "skewed append failed eventually, last_error={e:?}, shard_id={}",
                            shard.shard_id
                        );
                    }
                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                    continue;
                }
            };

            let status = resp.status;
            let done = match E2ETestSuite::check_append_response(shard.shard_id, status) {
                Ok(ok) => ok,
                Err(e) => panic!(
                    "skewed append critical error={e:?}, shard_id={}",
                    shard.shard_id
                ),
            };
            if !done {
                att += 1;
                {
                    let mut guard = retry_stats.lock().unwrap_or_else(|e| e.into_inner());
                    *guard.entry(status).or_insert(0) += 1;
                }
                if att > 10 {
                    panic!(
                        "skewed append can't append eventually, shard_id={}, last_status={}",
                        shard.shard_id, status
                    );
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                continue;
            }

            info!(
                "skewed append end {}/{} @ {}, elapsed={}ms, range={:?}",
                unit,
                op,
                shard.shard_id,
                SystemTime::now().duration_since(start).unwrap().as_millis(),
                append_req
                    .frag_metas
                    .first()
                    .and_then(|f| f.property.as_ref())
                    .map(|p| (p.min_handle.clone(), p.max_handle.clone()))
            );
            return shard.shard_id;
        }
    }

    pub async fn boot(&self, suite: Arc<RwLock<E2ETestSuite>>) -> Arc<RwLock<ShardInfo>> {
        suite
            .write()
            .await
            .must_create_index(&self.index_key)
            .await
            .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        let index = suite.read().await.get_index(&self.index_key).await;
        let shard = index.select_first_shard().await.unwrap();
        let addr = suite.read().await.get_one_worker_node_addr();
        {
            let mock = suite.read().await.get_mock_worker(&addr).await;
            mock.respond_compact_delay
                .store(500, std::sync::atomic::Ordering::Relaxed);
            mock.after_compaction_ratio
                .store(1.0, std::sync::atomic::Ordering::Relaxed);
            mock.mock_compact_process
                .store(true, std::sync::atomic::Ordering::Relaxed);
        }
        shard
    }

    #[cfg(target_os = "linux")]
    pub async fn stop_and_print_result(
        &self,
        suite: &mut E2ETestSuite,
        _addr: String,
        output: Option<&mut std::fs::File>,
    ) {
        use tici::ShardInfo;
        let worker_addrs = suite.workers.keys().cloned().collect::<Vec<_>>();
        let mut rows_from_frag = 0;
        let mut expected_rows_from_frag = 0;
        for addr in worker_addrs {
            let mock = suite.get_mock_worker(&addr).await;
            rows_from_frag += mock.get_total_rows_by_frags().await;
            expected_rows_from_frag += mock.get_total_expected_rows_by_frags().await;
        }
        let shard_info_heap_size = {
            use heapsize::HeapSizeOf;
            suite
                .collect_shards_from_index(&self.index_key)
                .await
                .iter()
                .map(|x| x.heap_size_of_children() + std::mem::size_of::<ShardInfo>())
                .sum::<usize>()
        };
        suite.stop_worker().await.unwrap();

        // Sleep some time, so the memory usage are stable.
        if std::env::var("TEST_BENCH_ROWS").is_ok() {
            tokio::time::sleep(std::time::Duration::from_millis(3000)).await;
        }

        {
            use std::io::Write;

            use log::warn;
            let elapsed = suite.start.elapsed().unwrap();

            let index = suite.get_index(&self.index_key).await;
            // only import log::warn here to make compiler happy
            if let Ok(p) = procfs::process::Process::myself() {
                // Every time we append, it will add `default_rows_batch` rows, which contribute
                // to 1 in `rows_from_frag`.
                assert_eq!(rows_from_frag, expected_rows_from_frag);
                assert_eq!(expected_rows_from_frag, self.n_default);
                let expected_rows = self.default_rows_batch * self.n_default;
                let m = p.stat.rss * unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
                let total_rows = suite.compute_total_rows(&self.index_key).await;
                let total_size =
                    ReadableSize(suite.compute_total_size(&self.index_key).await).as_mb_f64();
                let total_shards = suite.collect_shards_from_restore().await.len();
                let result = format!(
                    "Final result: workload={}, Mem={:.2} MB, Time={} s, total_size={:.2} MB, total_rows={total_rows}({:.4}) of expected {expected_rows}, total_frags={}, total_shards={}, shard_avg_size={:.2} MB, shard_heap_size={:.2} MB, index_stats({})\n",
                    self.to_debug_string(),
                    ReadableSize(m as u64).as_mb_f64(),
                    elapsed.as_secs(),
                    total_size,
                    total_rows as f64 / expected_rows as f64,
                    suite.compute_total_frags(&self.index_key).await,
                    total_shards,
                    total_size / total_shards as f64 / 1024.0 / 1024.0,
                    ReadableSize(shard_info_heap_size as u64).as_mb_f64(),
                    index.stats.to_short_string(),
                );
                if let Some(s) = output {
                    s.write_all(result.as_bytes()).unwrap();
                }
                warn!("{}", result);
            }
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub async fn stop_and_print_result(
        &self,
        suite: &mut E2ETestSuite,
        _addr: String,
        _output: Option<&mut std::fs::File>,
    ) {
        // On non-Linux platforms, heap size calculation is not available
        suite.stop_worker().await.unwrap();

        // Sleep some time, so the memory usage are stable.
        if std::env::var("TEST_BENCH_ROWS").is_ok() {
            tokio::time::sleep(std::time::Duration::from_millis(3000)).await;
        }
    }

    pub fn print_heap_size(&self) {
        if std::env::var("TEST_BENCH_ROWS").is_ok() {
            let result =
                tici::status_server::profile::dump_one_heap_profile(false, "svg".to_string());
            match result {
                Ok(body) => {
                    let exe: String = std::env::current_exe()
                        .expect("error get exe path")
                        .to_string_lossy()
                        .into_owned();
                    info!("dump or get heap profile successfully, exe={:?}", exe);
                    use std::{fs::File, io::Write};
                    let heap = format!("/heap_{}_{}.heap", self.n_default, self.n_parallel);
                    let mut file = File::create(heap.clone()).expect("Failed to create svg");
                    file.write_all(&body)
                        .expect("Failed to write heap profile to file");
                    info!(
                        "./jeprof.in -svg {} {} > heap_{}_{}.svg",
                        exe, heap, self.n_default, self.n_parallel
                    )
                }
                Err(e) => {
                    info!("dump or get heap profile fail, err={}", e);
                }
            }
        }
    }
}