use anyhow::{Context, Result, ensure};
use serde::{
    Deserialize, Deserializer, Serialize,
    de::{self, Unexpected, Visitor},
};
use tici_common::{
    config::{LogConfig, S3Config, ServerConfig},
    security::SecurityConfig,
};
use tici_system::{
    SysQuota,
    readable::{ReadableDuration, ReadableSize},
};

use crate::reader::shard_reader::{
    DEFAULT_FRAG_SAMPLE_COUNT, DEFAULT_SHARD_SAMPLE_COUNT, EstimateSampleCounts,
};

// WorkerNodeConfig represents the ** top-level ** configuration for a
// tici-worker node.
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct WorkerNodeConfig {
    pub s3: S3Config,
    pub frag_writer: FragWriterConfig,
    pub server: ServerConfig,
    pub heartbeat_interval: ReadableDuration,
    pub security: SecurityConfig,
    pub logger: LogConfig,
    pub compaction: CompactionConfig,
    pub storage: StorageConfig,
    // Meta service wait timeout in seconds. -1 means wait forever.
    pub meta_service_wait_timeout: i32,
    pub frag_upload: FragBatchUploaderConfig,
}

impl Default for WorkerNodeConfig {
    fn default() -> Self {
        Self {
            s3: S3Config::default(),
            frag_writer: FragWriterConfig::default(),
            server: ServerConfig::default(),
            heartbeat_interval: ReadableDuration::secs(3),
            security: SecurityConfig::default(),
            logger: LogConfig::default(),
            compaction: CompactionConfig::default(),
            storage: StorageConfig::default(),
            meta_service_wait_timeout: -1,
            frag_upload: FragBatchUploaderConfig::default(),
        }
    }
}

impl WorkerNodeConfig {
    pub fn load(path: &str) -> Result<Self> {
        let content = std::fs::read_to_string(path)?;
        let config: WorkerNodeConfig = toml::from_str(&content)?;

        config.server.grpc.validate()?;
        config.server.grpc_client.validate()?;
        config.frag_writer.validate()?;

        Ok(config)
    }

    pub fn new_for_test(data_path: &str) -> Self {
        Self {
            storage: StorageConfig::new(data_path.to_string()),
            s3: S3Config::default_for_test(),
            ..Self::default()
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct FragWriterConfig {
    pub index_num_threads: usize,

    pub index_mem_budget: ReadableSize,

    pub index_flush_interval: ReadableDuration,

    pub index_flush_size_limit: ReadableSize,

    // The interval for polling new CDC files from S3.
    pub cdc_poll_interval: ReadableDuration,
    // The maximum interval between polling new CDC files when empty polls
    // keep happening consecutively.
    pub cdc_max_poll_interval: ReadableDuration,

    // The sender channel size from CDC file reader to each shard writer.
    pub cdc_sender_channel_size: usize,
    // Number of CDC files to prefetch concurrently while preserving sequential
    // processing order.
    pub cdc_prefetch_window: usize,
    // Whether to fallback to head_object when get_object metadata does not
    // include last_modified.
    pub cdc_head_object_last_modified_fallback: bool,
    // Maximum number of CDC list requests allowed to hit storage backend concurrently.
    pub cdc_list_max_concurrency: usize,
}

impl Default for FragWriterConfig {
    fn default() -> Self {
        FragWriterConfig {
            index_num_threads: 1,
            index_mem_budget: ReadableSize::mb(128),
            index_flush_interval: ReadableDuration::secs(3),
            index_flush_size_limit: ReadableSize::mb(32),
            cdc_poll_interval: ReadableDuration::secs(1),
            cdc_max_poll_interval: ReadableDuration::secs(1),
            cdc_sender_channel_size: 2048,
            cdc_prefetch_window: 10,
            cdc_head_object_last_modified_fallback: true,
            cdc_list_max_concurrency: 128,
        }
    }
}

impl FragWriterConfig {
    pub fn validate(&self) -> Result<()> {
        ensure!(
            self.cdc_prefetch_window > 0,
            "cdc_prefetch_window must be greater than 0"
        );
        ensure!(
            self.cdc_max_poll_interval > ReadableDuration::secs(0),
            "cdc_max_poll_interval must be greater than 0"
        );
        ensure!(
            self.cdc_max_poll_interval >= self.cdc_poll_interval,
            "cdc_max_poll_interval must be greater than or equal to cdc_poll_interval"
        );
        ensure!(
            self.cdc_list_max_concurrency > 0,
            "cdc_list_max_concurrency must be greater than 0"
        );
        Ok(())
    }
}

// CompactionConfig holds configuration related to data compaction and
// data import.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct CompactionConfig {
    pub max_concurrency: usize,
    pub index_num_threads: usize,
    pub index_mem_budget: ReadableSize,
}

impl Default for CompactionConfig {
    fn default() -> Self {
        Self {
            // TODO: decide based on CPU cores?
            max_concurrency: 40,
            index_num_threads: 1,
            index_mem_budget: ReadableSize::gb(1),
        }
    }
}

// StorageConfig holds configuration related to local storage paths.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct StorageConfig {
    pub data_dir: String,
}

impl StorageConfig {
    pub fn new(data_dir: String) -> Self {
        Self { data_dir }
    }

    pub fn overwrite_by_cmdline(&mut self, v: &Option<String>) {
        if let Some(data_dir) = v {
            if !data_dir.is_empty() {
                self.data_dir = data_dir.to_string()
            }
        }
    }

    pub fn frag_base_path(&self) -> String {
        format!("{}/fragments", self.data_dir)
    }

    pub fn metadata_path(&self) -> String {
        format!("{}/metadata", self.data_dir)
    }

    pub fn temporary_path(&self) -> String {
        format!("{}/temp", self.data_dir)
    }

    pub fn init_worker_data_dir(&self) -> Result<()> {
        // The completed frag data files will be uploaded to S3 and then deleted. Upon
        // process startup, residual files in the directory can be directly removed.
        let frag_base_path = self.frag_base_path();
        if std::path::Path::new(&frag_base_path).exists() {
            std::fs::remove_dir_all(&frag_base_path)
                .with_context(|| format!("remove_dir_all: path={}", frag_base_path))?;
        }
        std::fs::create_dir_all(&frag_base_path)
            .with_context(|| format!("create_dir_all: path={}", frag_base_path))?;
        Ok(())
    }

    pub fn init_reader_data_dir(&self) -> Result<()> {
        let frag_base_path = self.frag_base_path();
        if !std::path::Path::new(&frag_base_path).exists() {
            std::fs::create_dir_all(&frag_base_path)
                .with_context(|| format!("create_dir_all: path={}", frag_base_path))?;
        }
        Ok(())
    }
}

impl Default for StorageConfig {
    fn default() -> Self {
        Self {
            data_dir: "tici-worker-data".to_string(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct FragReaderConfig {
    // Process-level decoded document cache capacity shared by all frag readers.
    // Bare TOML floats, or unitless strings in 0.xxx form, are parsed as memory
    // quota ratios up to 0.8. Integers and size strings are parsed as ReadableSize.
    #[serde(deserialize_with = "deserialize_global_doc_cache_size")]
    pub global_doc_cache_size: ReadableSize,
    pub remove_zero_ref_frag_delay: ReadableDuration,
    pub disk_reserve_size: ReadableSize,
    pub shard_sample_count: usize,
    pub frag_sample_count: usize,
}

const DEFAULT_GLOBAL_DOC_CACHE_MEMORY_DIVISOR: u64 = 10;
const MAX_GLOBAL_DOC_CACHE_MEMORY_RATIO: f64 = 0.8;

impl FragReaderConfig {
    pub fn estimate_sample_counts(&self) -> EstimateSampleCounts {
        EstimateSampleCounts {
            shard_sample_count: self.shard_sample_count,
            frag_sample_count: self.frag_sample_count,
        }
    }
}

pub fn default_global_doc_cache_size() -> ReadableSize {
    global_doc_cache_size_from_memory_limit(SysQuota::memory_limit_in_bytes())
}

fn global_doc_cache_size_from_memory_limit(memory_limit: u64) -> ReadableSize {
    ReadableSize(memory_limit / DEFAULT_GLOBAL_DOC_CACHE_MEMORY_DIVISOR)
}

fn global_doc_cache_size_from_memory_limit_ratio(
    memory_limit: u64,
    ratio: f64,
) -> Result<ReadableSize, String> {
    if !ratio.is_finite() || !(0.0..=MAX_GLOBAL_DOC_CACHE_MEMORY_RATIO).contains(&ratio) {
        return Err(format!(
            "global-doc-cache-size ratio must be between 0.0 and {MAX_GLOBAL_DOC_CACHE_MEMORY_RATIO}, got {ratio}"
        ));
    }

    Ok(ReadableSize((memory_limit as f64 * ratio) as u64))
}

fn global_doc_cache_size_from_ratio(ratio: f64) -> Result<ReadableSize, String> {
    global_doc_cache_size_from_memory_limit_ratio(SysQuota::memory_limit_in_bytes(), ratio)
}

fn is_global_doc_cache_ratio_str(value: &str) -> bool {
    let value = value.trim();
    let Some(fractional) = value.strip_prefix("0.") else {
        return false;
    };

    !fractional.is_empty() && fractional.chars().all(|ch| ch.is_ascii_digit())
}

fn parse_global_doc_cache_size_str(value: &str) -> Result<ReadableSize, String> {
    if is_global_doc_cache_ratio_str(value) {
        let ratio = value
            .trim()
            .parse::<f64>()
            .map_err(|_| format!("invalid global-doc-cache-size ratio: {value:?}"))?;
        return global_doc_cache_size_from_ratio(ratio);
    }

    value.parse::<ReadableSize>()
}

fn deserialize_global_doc_cache_size<'de, D>(deserializer: D) -> Result<ReadableSize, D::Error>
where
    D: Deserializer<'de>,
{
    struct GlobalDocCacheSizeVisitor;

    impl<'de> Visitor<'de> for GlobalDocCacheSizeVisitor {
        type Value = ReadableSize;

        fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            formatter.write_str("a readable size or a memory quota ratio from 0.0 to 0.8")
        }

        fn visit_i64<E>(self, size: i64) -> Result<ReadableSize, E>
        where
            E: de::Error,
        {
            if size >= 0 {
                Ok(ReadableSize(size as u64))
            } else {
                Err(E::invalid_value(Unexpected::Signed(size), &self))
            }
        }

        fn visit_u64<E>(self, size: u64) -> Result<ReadableSize, E>
        where
            E: de::Error,
        {
            Ok(ReadableSize(size))
        }

        fn visit_f64<E>(self, ratio: f64) -> Result<ReadableSize, E>
        where
            E: de::Error,
        {
            global_doc_cache_size_from_ratio(ratio).map_err(E::custom)
        }

        fn visit_str<E>(self, value: &str) -> Result<ReadableSize, E>
        where
            E: de::Error,
        {
            parse_global_doc_cache_size_str(value).map_err(E::custom)
        }
    }

    deserializer.deserialize_any(GlobalDocCacheSizeVisitor)
}

impl Default for FragReaderConfig {
    fn default() -> Self {
        Self {
            global_doc_cache_size: default_global_doc_cache_size(),
            remove_zero_ref_frag_delay: ReadableDuration::secs(600),
            disk_reserve_size: ReadableSize::gb(10),
            shard_sample_count: DEFAULT_SHARD_SAMPLE_COUNT,
            frag_sample_count: DEFAULT_FRAG_SAMPLE_COUNT,
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct FragBatchUploaderConfig {
    pub max_batch_size: ReadableSize,
    pub max_batch_delay: ReadableDuration,
    pub max_batch_concurrency: usize,
}

impl Default for FragBatchUploaderConfig {
    fn default() -> Self {
        Self {
            max_batch_size: ReadableSize::mb(256),
            max_batch_delay: ReadableDuration::secs(1),
            max_batch_concurrency: 64,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;

    use tempfile::NamedTempFile;
    use tici_test_utils::test_utils::init_logger_for_ut;

    use super::*;

    #[test]
    fn test_load_valid_config() {
        init_logger_for_ut();

        const VALID_CONFIG: &str = r#"
heartbeat-interval = "3s"
meta-service-wait-timeout = 120

[s3]
endpoint = "http://localhost:9000"
region = "us-east-1"
access-key = "minioadmin"
secret-key = "minioadmin"
use-path-style = true

[storage]
data-dir = "mock-worker-dir"

[frag-writer]
index-num-threads = 2
index-mem-budget = "2048MB"
index-flush-interval = "5s"
index-flush-size-limit = "1048576KB"
cdc-poll-interval = "2s"
cdc-max-poll-interval = "20s"
cdc-sender-channel-size = 1024
cdc-prefetch-window = 4
cdc-head-object-last-modified-fallback = false
cdc-list-max-concurrency = 64

[compaction]
max-concurrency = 20
index-num-threads = 2
index-mem-budget = "512MB"

[server]
addr = "127.0.0.1:8510"
pd-addr = "localhost:2379"

[security]
ca-path = "/path/to/ca.pem"
cert-path = "/path/to/cert.pem"
key-path = "/path/to/key.pem"
cert-allowed-cn = []

[logger]
filename = "tici_worker.log"
level = "debug"
size = "32MB"
count = 5

"#;

        let mut file = NamedTempFile::new().unwrap();
        file.write_all(VALID_CONFIG.as_bytes()).unwrap();

        let config = WorkerNodeConfig::load(file.path().to_str().unwrap())
            .expect("Failed to load valid config");

        log::info!("Loaded config from file: {:?}", config);

        // s3
        assert_eq!(config.s3.endpoint, "http://localhost:9000");
        assert_eq!(config.s3.region, "us-east-1");
        assert_eq!(config.s3.access_key, "minioadmin");
        assert_eq!(config.s3.secret_key, "minioadmin");
        assert!(config.s3.use_path_style);

        // storage
        assert_eq!(config.storage.data_dir, "mock-worker-dir");

        // frag_writer
        assert_eq!(config.frag_writer.index_num_threads, 2);
        assert_eq!(config.frag_writer.index_mem_budget, ReadableSize::mb(2048));
        assert_eq!(
            config.frag_writer.index_flush_interval,
            ReadableDuration::secs(5)
        );
        assert_eq!(
            config.frag_writer.index_flush_size_limit,
            ReadableSize::kb(1048576)
        );
        assert_eq!(
            config.frag_writer.cdc_poll_interval,
            ReadableDuration::secs(2)
        );
        assert_eq!(
            config.frag_writer.cdc_max_poll_interval,
            ReadableDuration::secs(20)
        );
        assert_eq!(config.frag_writer.cdc_sender_channel_size, 1024);
        assert_eq!(config.frag_writer.cdc_prefetch_window, 4);
        assert!(!config.frag_writer.cdc_head_object_last_modified_fallback);
        assert_eq!(config.frag_writer.cdc_list_max_concurrency, 64);

        // compaction
        assert_eq!(config.compaction.max_concurrency, 20);
        assert_eq!(config.compaction.index_num_threads, 2);
        assert_eq!(config.compaction.index_mem_budget, ReadableSize::mb(512));

        // heartbeat_interval
        assert_eq!(config.heartbeat_interval, ReadableDuration::secs(3));

        // server
        assert_eq!(config.server.addr, "127.0.0.1:8510");
        assert_eq!(config.server.pd_addr, "localhost:2379");

        // security
        assert_eq!(config.security.ca_path, "/path/to/ca.pem");
        assert_eq!(config.security.cert_path, "/path/to/cert.pem");
        assert_eq!(config.security.key_path, "/path/to/key.pem");

        // logger
        assert_eq!(config.logger.filename, "tici_worker.log");
        assert_eq!(config.logger.get_level_or(""), "debug");
        assert_eq!(config.logger.size, ReadableSize::mb(32));
        assert_eq!(config.logger.count, 5);
    }

    #[test]
    fn test_load_empty_config() {
        let mut file = NamedTempFile::new().unwrap();

        const EMPTY_CONFIG: &str = r#"
        "#;

        file.write_all(EMPTY_CONFIG.as_bytes()).unwrap();

        // When parsed from an empty config, all fields should take default values.
        let result = WorkerNodeConfig::load(file.path().to_str().unwrap());
        assert!(result.is_ok());

        let config = result.unwrap();
        assert_eq!(config.s3, S3Config::default());
        assert_eq!(config.frag_writer, FragWriterConfig::default());
        assert_eq!(config.frag_writer.cdc_prefetch_window, 10);
        assert!(config.frag_writer.cdc_head_object_last_modified_fallback);
        assert_eq!(
            config.frag_writer.cdc_max_poll_interval,
            ReadableDuration::secs(1)
        );
        assert_eq!(config.frag_writer.cdc_list_max_concurrency, 128);
        assert_eq!(config.server, ServerConfig::default());
        assert_eq!(config.heartbeat_interval, ReadableDuration::secs(3));
        assert_eq!(config.security, SecurityConfig::default());
        assert_eq!(config.logger, LogConfig::default());
        assert_eq!(config.compaction, CompactionConfig::default());
        assert_eq!(config.storage, StorageConfig::default());
        assert_eq!(config.meta_service_wait_timeout, -1);

        log::info!("Loaded config from file: {:?}", config);
    }

    #[test]
    fn test_nonexistent_file() {
        let result = WorkerNodeConfig::load("/path/that/does/not/exist.toml");
        assert!(result.is_err(), "Should fail on nonexistent file");
        assert!(
            result.unwrap_err().to_string().contains("No such file"),
            "Error should mention file not found"
        );
    }

    #[test]
    fn test_cmd_overwrite() {
        let mut cfg = StorageConfig::default();
        assert_eq!(cfg.data_dir, "tici-worker-data");
        cfg.overwrite_by_cmdline(&None);
        assert_eq!(cfg.data_dir, "tici-worker-data");
        cfg.overwrite_by_cmdline(&Some("".to_string()));
        assert_eq!(cfg.data_dir, "tici-worker-data");
        cfg.overwrite_by_cmdline(&Some("new_storage_dir".to_string()));
        assert_eq!(cfg.data_dir, "new_storage_dir");
    }

    #[test]
    fn test_global_doc_cache_default_scales_with_memory_limit() {
        assert_eq!(global_doc_cache_size_from_memory_limit(0), ReadableSize(0));
        assert_eq!(
            global_doc_cache_size_from_memory_limit(ReadableSize::gb(4).0),
            ReadableSize(ReadableSize::gb(4).0 / 10)
        );
        assert_eq!(
            global_doc_cache_size_from_memory_limit(ReadableSize::gb(256).0),
            ReadableSize(ReadableSize::gb(256).0 / 10)
        );
    }

    #[test]
    fn test_global_doc_cache_size_supports_size_and_ratio() {
        let config: FragReaderConfig = toml::from_str(r#"global-doc-cache-size = "2GB""#).unwrap();
        assert_eq!(config.global_doc_cache_size, ReadableSize::gb(2));

        let config: FragReaderConfig = toml::from_str(r#"global-doc-cache-size = 0.1"#).unwrap();
        assert_eq!(
            config.global_doc_cache_size,
            global_doc_cache_size_from_memory_limit_ratio(SysQuota::memory_limit_in_bytes(), 0.1,)
                .unwrap()
        );

        let config: FragReaderConfig = toml::from_str(r#"global-doc-cache-size = "0.1""#).unwrap();
        assert_eq!(
            config.global_doc_cache_size,
            global_doc_cache_size_from_memory_limit_ratio(SysQuota::memory_limit_in_bytes(), 0.1,)
                .unwrap()
        );

        let config: FragReaderConfig = toml::from_str(r#"global-doc-cache-size = 0.8"#).unwrap();
        assert_eq!(
            config.global_doc_cache_size,
            global_doc_cache_size_from_memory_limit_ratio(
                SysQuota::memory_limit_in_bytes(),
                MAX_GLOBAL_DOC_CACHE_MEMORY_RATIO,
            )
            .unwrap()
        );

        let config: FragReaderConfig = toml::from_str(r#"global-doc-cache-size = "0.8""#).unwrap();
        assert_eq!(
            config.global_doc_cache_size,
            global_doc_cache_size_from_memory_limit_ratio(
                SysQuota::memory_limit_in_bytes(),
                MAX_GLOBAL_DOC_CACHE_MEMORY_RATIO,
            )
            .unwrap()
        );
    }

    #[test]
    fn test_global_doc_cache_size_rejects_invalid_ratio() {
        let result = toml::from_str::<FragReaderConfig>(r#"global-doc-cache-size = 1.1"#);
        assert!(result.is_err());

        let result = toml::from_str::<FragReaderConfig>(r#"global-doc-cache-size = 0.81"#);
        assert!(result.is_err());

        let result = toml::from_str::<FragReaderConfig>(r#"global-doc-cache-size = "0.81""#);
        assert!(result.is_err());

        let result = toml::from_str::<FragReaderConfig>(r#"global-doc-cache-size = -0.1"#);
        assert!(result.is_err());
    }

    #[test]
    fn test_global_doc_cache_ratio_str_format() {
        assert!(is_global_doc_cache_ratio_str("0.1"));
        assert!(is_global_doc_cache_ratio_str("0.01"));
        assert!(is_global_doc_cache_ratio_str(" 0.001 "));

        // Unitless integers stay on the ReadableSize path, so "0" can disable
        // the cache as an absolute size.
        assert!(!is_global_doc_cache_ratio_str("0"));
        // A ratio string must include at least one digit after the decimal point.
        assert!(!is_global_doc_cache_ratio_str("0."));
        // Require the explicit 0.xxx form instead of accepting shorthand decimals.
        assert!(!is_global_doc_cache_ratio_str(".1"));
        // Ratio strings must be less than 1.0 syntactically; 1.0 is not 0.xxx.
        assert!(!is_global_doc_cache_ratio_str("1.0"));
        // Scientific notation is intentionally unsupported for ratio strings.
        assert!(!is_global_doc_cache_ratio_str("1e-1"));
        // Values with units are fixed sizes and should be parsed by ReadableSize.
        assert!(!is_global_doc_cache_ratio_str("0.1GB"));
    }

    #[test]
    fn test_global_doc_cache_size_omitted_matches_struct_default() {
        let config: FragReaderConfig = toml::from_str(
            r#"
            shard-sample-count = 6
            "#,
        )
        .unwrap();
        assert_eq!(config.shard_sample_count, 6);
        assert_eq!(
            config.global_doc_cache_size,
            default_global_doc_cache_size(),
        );
        assert_eq!(
            config.global_doc_cache_size,
            FragReaderConfig::default().global_doc_cache_size,
        );
    }
}