use std::net::SocketAddr;

use narsil_engine::{
    BatchingConfig, BurnBackend, BurnDlrmBackend, BurnMlpBackend, DlrmConfig, InferenceBackend,
    MlpConfig, WorkerPoolConfig, serve, serve_with_batching, serve_with_worker_pool,
};

fn env_f32(name: &str, default: f32) -> narsil_engine::Result<f32> {
    match std::env::var(name) {
        Ok(value) => value.parse().map_err(|err| {
            narsil_engine::NarsilError::InvalidRequest(format!(
                "invalid {name} value '{value}': {err}"
            ))
        }),
        Err(std::env::VarError::NotPresent) => Ok(default),
        Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "invalid {name}: {err}"
        ))),
    }
}

fn env_usize(name: &str, default: usize) -> narsil_engine::Result<usize> {
    match std::env::var(name) {
        Ok(value) => value.parse().map_err(|err| {
            narsil_engine::NarsilError::InvalidRequest(format!(
                "invalid {name} value '{value}': {err}"
            ))
        }),
        Err(std::env::VarError::NotPresent) => Ok(default),
        Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "invalid {name}: {err}"
        ))),
    }
}

fn env_usize_list(name: &str, default: &[usize]) -> narsil_engine::Result<Vec<usize>> {
    match std::env::var(name) {
        Ok(value) => value
            .split(',')
            .filter(|part| !part.trim().is_empty())
            .map(|part| {
                part.trim().parse().map_err(|err| {
                    narsil_engine::NarsilError::InvalidRequest(format!(
                        "invalid {name} value '{value}': {err}"
                    ))
                })
            })
            .collect(),
        Err(std::env::VarError::NotPresent) => Ok(default.to_vec()),
        Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "invalid {name}: {err}"
        ))),
    }
}

#[cfg(feature = "torch")]
fn env_string_list(name: &str, default: &[&str]) -> narsil_engine::Result<Vec<String>> {
    match std::env::var(name) {
        Ok(value) => Ok(value
            .split(',')
            .map(|part| part.trim().to_string())
            .filter(|part| !part.is_empty())
            .collect()),
        Err(std::env::VarError::NotPresent) => {
            Ok(default.iter().map(|part| part.to_string()).collect())
        }
        Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "invalid {name}: {err}"
        ))),
    }
}

fn env_path(name: &str) -> Option<std::path::PathBuf> {
    std::env::var_os(name).map(std::path::PathBuf::from)
}

#[cfg(feature = "cuda")]
fn dlrm_interaction_from_env() -> narsil_engine::Result<String> {
    match std::env::var("NARSIL_DLRM_INTERACTION") {
        Ok(value) => match value.as_str() {
            "burn" | "default" | "cutile" => Ok(value),
            other => Err(narsil_engine::NarsilError::InvalidRequest(format!(
                "unknown NARSIL_DLRM_INTERACTION '{other}'"
            ))),
        },
        Err(std::env::VarError::NotPresent) => Ok("burn".to_string()),
        Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "invalid NARSIL_DLRM_INTERACTION: {err}"
        ))),
    }
}

fn mlp_config_from_env() -> narsil_engine::Result<MlpConfig> {
    let defaults = MlpConfig::default();
    MlpConfig::new(
        env_usize("NARSIL_MLP_INPUT_DIM", defaults.input_dim)?,
        env_usize("NARSIL_MLP_HIDDEN_DIM", defaults.hidden_dim)?,
        env_usize("NARSIL_MLP_OUTPUT_DIM", defaults.output_dim)?,
    )
}

fn dlrm_config_from_env() -> narsil_engine::Result<DlrmConfig> {
    let defaults = DlrmConfig::default();
    let config = DlrmConfig {
        dense_in_features: env_usize("NARSIL_DLRM_DENSE_IN_FEATURES", defaults.dense_in_features)?,
        sparse_features: env_usize("NARSIL_DLRM_SPARSE_FEATURES", defaults.sparse_features)?,
        num_embeddings: env_usize("NARSIL_DLRM_NUM_EMBEDDINGS", defaults.num_embeddings)?,
        embedding_dim: env_usize("NARSIL_DLRM_EMBEDDING_DIM", defaults.embedding_dim)?,
        dense_arch_layer_sizes: env_usize_list(
            "NARSIL_DLRM_DENSE_ARCH",
            &defaults.dense_arch_layer_sizes,
        )?,
        over_arch_layer_sizes: env_usize_list(
            "NARSIL_DLRM_OVER_ARCH",
            &defaults.over_arch_layer_sizes,
        )?,
    };
    config.validate()?;
    Ok(config)
}

fn default_workers(cuda: bool) -> usize {
    if cuda {
        1
    } else {
        std::thread::available_parallelism()
            .map(|value| value.get().clamp(1, 4))
            .unwrap_or(1)
    }
}

fn worker_pool_config(cuda: bool) -> narsil_engine::Result<WorkerPoolConfig> {
    WorkerPoolConfig::new(
        env_usize("NARSIL_INFERENCE_WORKERS", default_workers(cuda))?,
        env_usize("NARSIL_INFERENCE_QUEUE_DEPTH", 1024)?,
    )
}

fn batching_config(cuda: bool) -> narsil_engine::Result<BatchingConfig> {
    BatchingConfig::new(
        env_usize("NARSIL_BATCH_SIZE", 32)?,
        env_usize("NARSIL_BATCH_TIMEOUT_MS", 5)? as u64,
        env_usize("NARSIL_INFERENCE_WORKERS", default_workers(cuda))?,
        env_usize("NARSIL_INFERENCE_QUEUE_DEPTH", 1024)?,
    )
}

async fn serve_backend<B: InferenceBackend + Clone>(
    addr: SocketAddr,
    backend: B,
    cuda: bool,
) -> narsil_engine::Result<()> {
    let execution = std::env::var("NARSIL_EXECUTION_MODE").unwrap_or_else(|_| "worker".to_string());
    match execution.as_str() {
        "inline" => serve(addr, backend).await,
        "worker" | "worker-pool" | "async" => {
            let config = worker_pool_config(cuda)?;
            tracing::info!(?config, execution, "using async inference worker pool");
            serve_with_worker_pool(addr, backend, config).await
        }
        "batch" => {
            let config = batching_config(cuda)?;
            tracing::info!(?config, execution, "using batch accumulation mode");
            serve_with_batching(addr, backend, config).await
        }
        other => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "unknown NARSIL_EXECUTION_MODE '{other}'"
        ))),
    }
}

#[tokio::main]
async fn main() -> narsil_engine::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "narsil=info".into()),
        )
        .init();

    let addr: SocketAddr = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "127.0.0.1:50051".to_string())
        .parse()
        .map_err(|err| {
            narsil_engine::NarsilError::InvalidRequest(format!("invalid bind addr: {err}"))
        })?;

    let backend = std::env::var("NARSIL_BACKEND").unwrap_or_else(|_| "burn-affine".to_string());

    match backend.as_str() {
        "burn-affine" | "affine" => {
            let scale = env_f32("NARSIL_BURN_SCALE", 1.0)?;
            let bias = env_f32("NARSIL_BURN_BIAS", 0.0)?;
            tracing::info!(%addr, scale, bias, backend, "starting Narsil gRPC inference server");
            serve_backend(addr, BurnBackend::with_affine(scale, bias), false).await
        }
        "burn-mlp" | "mlp" => {
            let config = mlp_config_from_env()?;
            tracing::info!(%addr, ?config, backend, "starting Narsil gRPC inference server");
            if let Some(path) = env_path("NARSIL_MODEL_PATH") {
                tracing::info!(?path, "loading model weights from file");
                serve_backend(addr, BurnMlpBackend::flex_load(&path, config)?, false).await
            } else {
                serve_backend(addr, BurnMlpBackend::flex(config)?, false).await
            }
        }
        "burn-dlrm" | "dlrm" => {
            let config = dlrm_config_from_env()?;
            tracing::info!(%addr, ?config, backend, "starting Narsil gRPC inference server");
            if let Some(path) = env_path("NARSIL_MODEL_PATH") {
                tracing::info!(?path, "loading DLRM weights from file");
                serve_backend(addr, BurnDlrmBackend::flex_load(&path, config)?, false).await
            } else {
                serve_backend(addr, BurnDlrmBackend::flex(config)?, false).await
            }
        }
        "burn-cuda-mlp" | "cuda-mlp" => serve_cuda_mlp(addr).await,
        "burn-cuda-dlrm" | "cuda-dlrm" => serve_cuda_dlrm(addr).await,
        "burn-cuda-dlrm-fp16" | "cuda-dlrm-fp16" | "burn-cuda-dlrm-f16" | "cuda-dlrm-f16" => {
            serve_cuda_dlrm_with_precision(addr, CudaDlrmPrecision::Fp16).await
        }
        "torch-dlrm" | "torch" => serve_torch_dlrm(addr).await,
        other => Err(narsil_engine::NarsilError::InvalidRequest(format!(
            "unknown NARSIL_BACKEND '{other}'"
        ))),
    }
}

#[cfg(feature = "cuda")]
#[derive(Debug, Clone, Copy)]
enum CudaDlrmPrecision {
    Fp32,
    Fp16,
}

#[cfg(feature = "cuda")]
impl CudaDlrmPrecision {
    fn from_env() -> narsil_engine::Result<Self> {
        match std::env::var("NARSIL_DLRM_PRECISION") {
            Ok(value) => match value.as_str() {
                "fp32" | "f32" | "float32" => Ok(Self::Fp32),
                "fp16" | "f16" | "float16" | "half" => Ok(Self::Fp16),
                other => Err(narsil_engine::NarsilError::InvalidRequest(format!(
                    "unknown NARSIL_DLRM_PRECISION '{other}'"
                ))),
            },
            Err(std::env::VarError::NotPresent) => Ok(Self::Fp32),
            Err(err) => Err(narsil_engine::NarsilError::InvalidRequest(format!(
                "invalid NARSIL_DLRM_PRECISION: {err}"
            ))),
        }
    }
}

#[cfg(not(feature = "cuda"))]
#[derive(Debug, Clone, Copy)]
enum CudaDlrmPrecision {
    Fp16,
}

#[cfg(feature = "cuda")]
async fn serve_cuda_mlp(addr: SocketAddr) -> narsil_engine::Result<()> {
    let config = mlp_config_from_env()?;
    let device = burn_cuda::CudaDevice::new(env_usize("NARSIL_CUDA_DEVICE", 0)?);
    tracing::info!(%addr, ?config, ?device, "starting Narsil CUDA gRPC inference server");
    serve_backend(
        addr,
        BurnMlpBackend::<burn_cuda::Cuda<f32, i32>>::new(
            config,
            device,
            "burn-cuda-mlp",
            "burn-cuda",
        )?,
        true,
    )
    .await
}

#[cfg(feature = "cuda")]
async fn serve_cuda_dlrm(addr: SocketAddr) -> narsil_engine::Result<()> {
    serve_cuda_dlrm_with_precision(addr, CudaDlrmPrecision::from_env()?).await
}

#[cfg(feature = "cuda")]
async fn serve_cuda_dlrm_with_precision(
    addr: SocketAddr,
    precision: CudaDlrmPrecision,
) -> narsil_engine::Result<()> {
    let config = dlrm_config_from_env()?;
    let device = burn_cuda::CudaDevice::new(env_usize("NARSIL_CUDA_DEVICE", 0)?);
    let interaction = dlrm_interaction_from_env()?;
    tracing::info!(%addr, ?config, ?device, ?precision, %interaction, "starting Narsil CUDA DLRM gRPC inference server");
    match precision {
        CudaDlrmPrecision::Fp32 if interaction == "cutile" => {
            serve_cuda_dlrm_cutile_fp32(addr, config, device).await
        }
        CudaDlrmPrecision::Fp32 => {
            if let Some(path) = env_path("NARSIL_MODEL_PATH") {
                tracing::info!(?path, "loading DLRM weights from file");
                serve_backend(
                    addr,
                    BurnDlrmBackend::<burn_cuda::Cuda<f32, i32>>::load(
                        &path,
                        config,
                        device,
                        "burn-cuda-dlrm-fp32",
                        "burn-cuda-fp32",
                    )?,
                    true,
                )
                .await
            } else {
                serve_backend(
                    addr,
                    BurnDlrmBackend::<burn_cuda::Cuda<f32, i32>>::new(
                        config,
                        device,
                        "burn-cuda-dlrm-fp32",
                        "burn-cuda-fp32",
                    )?,
                    true,
                )
                .await
            }
        }
        CudaDlrmPrecision::Fp16 => {
            if interaction == "cutile" {
                return Err(narsil_engine::NarsilError::InvalidRequest(
                    "NARSIL_DLRM_INTERACTION=cutile currently supports only NARSIL_DLRM_PRECISION=fp32"
                        .to_string(),
                ));
            }
            if let Some(path) = env_path("NARSIL_MODEL_PATH") {
                tracing::info!(?path, "loading DLRM weights from file");
                serve_backend(
                    addr,
                    BurnDlrmBackend::<burn_cuda::Cuda<burn_tensor::f16, i32>>::load(
                        &path,
                        config,
                        device,
                        "burn-cuda-dlrm-fp16",
                        "burn-cuda-fp16",
                    )?,
                    true,
                )
                .await
            } else {
                serve_backend(
                    addr,
                    BurnDlrmBackend::<burn_cuda::Cuda<burn_tensor::f16, i32>>::new(
                        config,
                        device,
                        "burn-cuda-dlrm-fp16",
                        "burn-cuda-fp16",
                    )?,
                    true,
                )
                .await
            }
        }
    }
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
async fn serve_cuda_dlrm_cutile_fp32(
    addr: SocketAddr,
    config: DlrmConfig,
    device: burn_cuda::CudaDevice,
) -> narsil_engine::Result<()> {
    if let Some(path) = env_path("NARSIL_MODEL_PATH") {
        tracing::info!(?path, "loading DLRM weights from file");
        serve_backend(
            addr,
            narsil_engine::CutileDlrmBackend::load(
                &path,
                config,
                device,
                "burn-cuda-dlrm-fp32-cutile",
                "burn-cuda-fp32-cutile",
            )?,
            true,
        )
        .await
    } else {
        serve_backend(
            addr,
            narsil_engine::CutileDlrmBackend::new(
                config,
                device,
                "burn-cuda-dlrm-fp32-cutile",
                "burn-cuda-fp32-cutile",
            )?,
            true,
        )
        .await
    }
}

#[cfg(all(feature = "cuda", not(feature = "cuda-cutile")))]
async fn serve_cuda_dlrm_cutile_fp32(
    _addr: SocketAddr,
    _config: DlrmConfig,
    _device: burn_cuda::CudaDevice,
) -> narsil_engine::Result<()> {
    Err(narsil_engine::NarsilError::InvalidRequest(
        "NARSIL_DLRM_INTERACTION=cutile requires building with --features cuda-cutile".to_string(),
    ))
}

#[cfg(not(feature = "cuda"))]
async fn serve_cuda_dlrm(_addr: SocketAddr) -> narsil_engine::Result<()> {
    Err(narsil_engine::NarsilError::InvalidRequest(
        "NARSIL_BACKEND=burn-cuda-dlrm requires building with --features cuda".to_string(),
    ))
}

#[cfg(not(feature = "cuda"))]
async fn serve_cuda_dlrm_with_precision(
    _addr: SocketAddr,
    _precision: CudaDlrmPrecision,
) -> narsil_engine::Result<()> {
    Err(narsil_engine::NarsilError::InvalidRequest(
        "NARSIL_BACKEND=burn-cuda-dlrm-fp16 requires building with --features cuda".to_string(),
    ))
}

#[cfg(not(feature = "cuda"))]
async fn serve_cuda_mlp(_addr: SocketAddr) -> narsil_engine::Result<()> {
    Err(narsil_engine::NarsilError::InvalidRequest(
        "NARSIL_BACKEND=burn-cuda-mlp requires building with --features cuda".to_string(),
    ))
}

#[cfg(feature = "torch")]
async fn serve_torch_dlrm(addr: SocketAddr) -> narsil_engine::Result<()> {
    let config = dlrm_config_from_env()?;
    // The FBGEMM INT8 inference op is GPU-only, so this backend requires a CUDA device.
    if matches!(std::env::var("NARSIL_CUDA_DEVICE").as_deref(), Ok("cpu")) {
        return Err(narsil_engine::NarsilError::InvalidRequest(
            "NARSIL_BACKEND=torch-dlrm requires a CUDA device; CPU is not supported for the \
             FBGEMM inference path"
                .to_string(),
        ));
    }
    let device_index = env_usize("NARSIL_CUDA_DEVICE", 0)?;
    let model_path = env_path("NARSIL_TORCH_MODEL_PATH").ok_or_else(|| {
        narsil_engine::NarsilError::InvalidRequest(
            "NARSIL_BACKEND=torch-dlrm requires NARSIL_TORCH_MODEL_PATH".to_string(),
        )
    })?;
    // Ordered list of fbgemm .so plugins to dlopen before loading; these register fbgemm:: ops.
    let preload_libs = env_string_list("NARSIL_TORCH_PRELOAD_LIBS", &[])?;
    if preload_libs.is_empty() {
        tracing::warn!(
            "NARSIL_TORCH_PRELOAD_LIBS is empty; fbgemm ops will likely fail. \
             Set it to the fbgemm inference .so set (comma-separated)."
        );
    }
    tracing::info!(%addr, ?config, device_index, ?model_path, "starting Narsil torch DLRM gRPC inference server");
    let backend = narsil_engine::TorchDlrmBackend::load(
        &model_path,
        config,
        device_index,
        &preload_libs,
        "torch-dlrm",
        "libtorch-fbgemm",
    )?;
    serve_backend(addr, backend, true).await
}

#[cfg(not(feature = "torch"))]
async fn serve_torch_dlrm(_addr: SocketAddr) -> narsil_engine::Result<()> {
    Err(narsil_engine::NarsilError::InvalidRequest(
        "NARSIL_BACKEND=torch-dlrm requires building with --features torch".to_string(),
    ))
}
