cgpu 0.1.0

use std::fmt;
use std::fs;
use std::path::Path;
use std::sync::{mpsc, Arc};

use serde::Deserialize;

use crate::error::ExecutorInitError;
use crate::telemetry::TelemetryReport;

pub const DEFAULT_CONFIG_PATH: &str = "cgpu.config.json";
pub const DEFAULT_MEMORY_FILL_RATIO: f32 = 0.90;
pub const DEFAULT_MIN_BATCH_BYTES: usize = 4096;
pub const DEFAULT_MAX_BATCH_BYTES: usize = 256 * 1024 * 1024;

#[derive(Clone, Debug)]
pub struct ExecutorConfig {
    pub vram_override: Option<usize>,
    pub memory_fill_ratio: f32,
    pub min_batch_bytes: usize,
    pub max_batch_bytes: usize,
    pub execution_mode: ExecutionMode,
    pub cpu_threads: Option<usize>,
    pub parallel_fallback: bool,
    pub shader_cache: bool,
    pub shader_optimization: ShaderOpt,
    pub enable_telemetry: bool,
    pub telemetry_sink: TelemetrySink,
}

impl Default for ExecutorConfig {
    fn default() -> Self {
        Self {
            vram_override: None,
            memory_fill_ratio: DEFAULT_MEMORY_FILL_RATIO,
            min_batch_bytes: DEFAULT_MIN_BATCH_BYTES,
            max_batch_bytes: DEFAULT_MAX_BATCH_BYTES,
            execution_mode: ExecutionMode::Auto,
            cpu_threads: None,
            parallel_fallback: true,
            shader_cache: true,
            shader_optimization: ShaderOpt::Performance,
            enable_telemetry: true,
            telemetry_sink: TelemetrySink::Log,
        }
    }
}

impl ExecutorConfig {
    pub fn from_json_file(path: impl AsRef<Path>) -> Result<Self, ExecutorInitError> {
        let text = fs::read_to_string(path.as_ref()).map_err(|error| {
            ExecutorInitError::InvalidConfig(format!(
                "failed to read {}: {error}",
                path.as_ref().display()
            ))
        })?;
        let file_config: ExecutorConfigFile = serde_json::from_str(&text).map_err(|error| {
            ExecutorInitError::InvalidConfig(format!(
                "failed to parse {}: {error}",
                path.as_ref().display()
            ))
        })?;
        Ok(file_config.apply_to(Self::default()))
    }

    pub fn from_default_json_file() -> Result<Self, ExecutorInitError> {
        for path in default_config_paths() {
            if path.exists() {
                return Self::from_json_file(path);
            }
        }
        Self::from_json_file(DEFAULT_CONFIG_PATH)
    }

    pub fn from_default_json_file_or_default() -> Result<Self, ExecutorInitError> {
        if default_config_paths().into_iter().any(|path| path.exists()) {
            Self::from_default_json_file()
        } else {
            Ok(Self::default())
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ExecutionMode {
    Auto,
    PreferGpu,
    PreferCpu,
    GpuOnly,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ShaderOpt {
    Debug,
    Performance,
    Size,
}

#[derive(Clone)]
pub enum TelemetrySink {
    Log,
    Callback(Arc<dyn Fn(&TelemetryReport) + Send + Sync + 'static>),
    Buffer(mpsc::Sender<TelemetryReport>),
}

impl fmt::Debug for TelemetrySink {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Log => f.write_str("Log"),
            Self::Callback(_) => f.write_str("Callback(..)"),
            Self::Buffer(_) => f.write_str("Buffer(..)"),
        }
    }
}

impl Default for TelemetrySink {
    fn default() -> Self {
        Self::Log
    }
}

#[derive(Debug, Default, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct ExecutorConfigFile {
    vram_override: Option<usize>,
    memory_fill_ratio: Option<f32>,
    min_batch_bytes: Option<usize>,
    max_batch_bytes: Option<usize>,
    execution_mode: Option<ExecutionModeConfig>,
    cpu_threads: Option<usize>,
    parallel_fallback: Option<bool>,
    shader_cache: Option<bool>,
    shader_optimization: Option<ShaderOptConfig>,
    enable_telemetry: Option<bool>,
    telemetry_sink: Option<TelemetrySinkConfig>,
}

impl ExecutorConfigFile {
    fn apply_to(self, mut config: ExecutorConfig) -> ExecutorConfig {
        if let Some(value) = self.vram_override {
            config.vram_override = Some(value);
        }
        if let Some(value) = self.memory_fill_ratio {
            config.memory_fill_ratio = value;
        }
        if let Some(value) = self.min_batch_bytes {
            config.min_batch_bytes = value;
        }
        if let Some(value) = self.max_batch_bytes {
            config.max_batch_bytes = value;
        }
        if let Some(value) = self.execution_mode {
            config.execution_mode = value.into();
        }
        if let Some(value) = self.cpu_threads {
            config.cpu_threads = Some(value);
        }
        if let Some(value) = self.parallel_fallback {
            config.parallel_fallback = value;
        }
        if let Some(value) = self.shader_cache {
            config.shader_cache = value;
        }
        if let Some(value) = self.shader_optimization {
            config.shader_optimization = value.into();
        }
        if let Some(value) = self.enable_telemetry {
            config.enable_telemetry = value;
        }
        if let Some(TelemetrySinkConfig::Log) = self.telemetry_sink {
            config.telemetry_sink = TelemetrySink::Log;
        }
        config
    }
}

#[derive(Clone, Copy, Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
enum ExecutionModeConfig {
    Auto,
    PreferGpu,
    PreferCpu,
    GpuOnly,
}

impl From<ExecutionModeConfig> for ExecutionMode {
    fn from(value: ExecutionModeConfig) -> Self {
        match value {
            ExecutionModeConfig::Auto => Self::Auto,
            ExecutionModeConfig::PreferGpu => Self::PreferGpu,
            ExecutionModeConfig::PreferCpu => Self::PreferCpu,
            ExecutionModeConfig::GpuOnly => Self::GpuOnly,
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
enum ShaderOptConfig {
    Debug,
    Performance,
    Size,
}

impl From<ShaderOptConfig> for ShaderOpt {
    fn from(value: ShaderOptConfig) -> Self {
        match value {
            ShaderOptConfig::Debug => Self::Debug,
            ShaderOptConfig::Performance => Self::Performance,
            ShaderOptConfig::Size => Self::Size,
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
enum TelemetrySinkConfig {
    Log,
}

fn default_config_paths() -> Vec<std::path::PathBuf> {
    let manifest_dir = Path::new(env!("CARGO_MANIFEST_DIR"));
    let mut paths = Vec::new();
    paths.push(Path::new(DEFAULT_CONFIG_PATH).to_path_buf());
    paths.push(manifest_dir.join(DEFAULT_CONFIG_PATH));
    if let Some(parent) = manifest_dir.parent() {
        paths.push(parent.join(DEFAULT_CONFIG_PATH));
    }
    paths
}