//! Native torch DLRM backend (route A): serve a TorchRec TorchScript artifact
//! through libtorch via tch-rs, so the real FBGEMM table-batched embedding and
//! interaction kernels run. The Narsil serving shell (tonic, worker pool, batch
//! collector) is unchanged — this is just another `InferenceBackend`.
//!
//! Contract (see TorchRec `dlrm_predict.py`): `forward(Dict[str, Tensor]) -> Dict[str, Tensor]`
//!   in : float_features f32[B,dense], id_list_features.{lengths,values} i32[B*sparse]
//!   out: "default" f32[B] (already on CPU)

use std::collections::HashMap;
use std::path::Path;

use libloading::os::unix::{Library, RTLD_GLOBAL, RTLD_NOW};
use tch::{CModule, Device, IValue, Kind, Tensor as TchTensor};

use crate::backend::dlrm::{DlrmConfig, find_tensor, int_values};
use crate::backend::{BackendMetadata, InferenceBackend};
use crate::error::{NarsilError, Result};
use crate::schema::{InferenceRequest, InferenceResponse, Tensor};

const DENSE_KEY: &str = "float_features";
const LENGTHS_KEY: &str = "id_list_features.lengths";
const VALUES_KEY: &str = "id_list_features.values";
const OUTPUT_KEY: &str = "default";

/// dlopen the fbgemm op libraries with RTLD_GLOBAL so the `fbgemm::*` ops the
/// artifact references resolve. Must run before module load.
/// Handles are leaked on purpose so the libraries stay resident.
fn preload_libraries(libs: &[String]) -> Result<()> {
    for path in libs {
        if !should_dlopen_preload(path) {
            continue;
        }
        // SAFETY: loading trusted libtorch/fbgemm shared objects; handle leaked deliberately.
        let lib = unsafe { Library::open(Some(path.as_str()), RTLD_NOW | RTLD_GLOBAL) }
            .map_err(|err| NarsilError::Backend(format!("failed to dlopen {path}: {err}")))?;
        std::mem::forget(lib);
    }
    Ok(())
}

#[cfg(narsil_libtorch_cuda_linked)]
fn should_dlopen_preload(path: &str) -> bool {
    if Path::new(path)
        .file_name()
        .is_some_and(|name| name == "libtorch_cuda.so")
    {
        return false;
    }
    true
}

#[cfg(not(narsil_libtorch_cuda_linked))]
fn should_dlopen_preload(_path: &str) -> bool {
    true
}

/// Checked i64 -> i32 narrowing for the sparse ids/lengths the TorchRec artifact consumes as int32.
fn to_i32(values: &[i64]) -> Result<Vec<i32>> {
    values
        .iter()
        .map(|&value| {
            i32::try_from(value).map_err(|_| {
                NarsilError::InvalidRequest(format!(
                    "DLRM sparse id/length {value} exceeds i32 range"
                ))
            })
        })
        .collect()
}

/// Reorder one or more **batch-major** requests into a single **key-major** (feature-major)
/// `KeyedJaggedTensor` layout: for each feature, every sample across every request is contiguous
/// (`[f0: req0.s0..req0.sN, req1.s0.., ...] [f1: ...] ...`). This is the layout TorchRec builds its
/// KJT from; concatenating batch-major blocks directly would scramble features across the batch.
/// Returns `(dense, lengths_key_major, values_key_major, total_batch)`; dense and the model output
/// stay in request-then-sample order, so per-request output slicing is unaffected.
fn assemble_key_major(
    prepared: &[&PreparedTorchInput],
    sparse_features: usize,
) -> (Vec<f32>, Vec<i64>, Vec<i64>, usize) {
    let total_batch: usize = prepared.iter().map(|p| p.batch).sum();

    let mut dense = Vec::new();
    for p in prepared {
        dense.extend_from_slice(&p.dense_values);
    }

    // Prefix offsets into each request's jagged batch-major `values`, indexed by `(b * F + f)`.
    let offsets: Vec<Vec<usize>> = prepared
        .iter()
        .map(|p| {
            let mut off = Vec::with_capacity(p.lengths.len() + 1);
            let mut acc = 0usize;
            off.push(0);
            for &len in &p.lengths {
                acc += len.max(0) as usize;
                off.push(acc);
            }
            off
        })
        .collect();

    let mut lengths_km = Vec::with_capacity(total_batch * sparse_features);
    let mut values_km = Vec::new();
    for f in 0..sparse_features {
        for (request_idx, p) in prepared.iter().enumerate() {
            for b in 0..p.batch {
                let cell = b * sparse_features + f;
                lengths_km.push(p.lengths[cell]);
                let (start, end) = (offsets[request_idx][cell], offsets[request_idx][cell + 1]);
                values_km.extend_from_slice(&p.values[start..end]);
            }
        }
    }

    (dense, lengths_km, values_km, total_batch)
}

struct TorchDlrmInner {
    module: CModule,
    device: Device,
    config: DlrmConfig,
    name: &'static str,
    runtime: &'static str,
}

/// A prepared, host-side request in **batch-major** layout (matching the Burn backend and the
/// Narsil wire format): `lengths[b * F + f]` is the id count for sample `b`, feature `f`, and
/// `values` is the jagged concatenation of ids in that same `(sample, feature)` order.
/// [`assemble_key_major`] reorders this into the key-major layout TorchRec's `KeyedJaggedTensor`
/// expects before it reaches the model.
struct PreparedTorchInput {
    dense_values: Vec<f32>,
    lengths: Vec<i64>,
    values: Vec<i64>,
    batch: usize,
}

#[derive(Clone)]
pub struct TorchDlrmBackend {
    inner: std::sync::Arc<TorchDlrmInner>,
}

impl TorchDlrmBackend {
    /// Load a TorchScript DLRM artifact onto a CUDA device after preloading the
    /// fbgemm shared libraries listed in `preload_libs`.
    pub fn load(
        model_path: &Path,
        config: DlrmConfig,
        device_index: usize,
        preload_libs: &[String],
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        config.validate()?;
        preload_libraries(preload_libs)?;

        let device = Device::Cuda(device_index);
        let module = CModule::load_on_device(model_path, device).map_err(|err| {
            NarsilError::Backend(format!(
                "failed to load torch DLRM module from {}: {err}",
                model_path.display()
            ))
        })?;

        Ok(Self {
            inner: std::sync::Arc::new(TorchDlrmInner {
                module,
                device,
                config,
                name,
                runtime,
            }),
        })
    }

    fn prepare_request(&self, request: &InferenceRequest) -> Result<PreparedTorchInput> {
        let dense = find_tensor(request, &["dense", DENSE_KEY])?;
        let values = find_tensor(request, &["values", VALUES_KEY])?;
        let lengths = find_tensor(request, &["lengths", LENGTHS_KEY])?;

        let dense_in = self.inner.config.dense_in_features;
        let batch = match dense.shape.as_slice() {
            [features] if *features == dense_in => 1,
            [batch, features] if *features == dense_in => *batch,
            _ => {
                return Err(NarsilError::InvalidRequest(format!(
                    "DLRM dense input expects shape [{dense_in}] or [batch, {dense_in}], got {:?}",
                    dense.shape
                )));
            }
        };

        let dense_values = dense.f32_values()?;
        let lengths = int_values(lengths)?;
        let values = int_values(values)?;

        // `lengths` is batch-major: one entry per (sample, feature).
        let expected = batch * self.inner.config.sparse_features;
        if lengths.len() != expected {
            return Err(NarsilError::InvalidRequest(format!(
                "DLRM expected {expected} sparse lengths (batch * sparse_features), got {}",
                lengths.len()
            )));
        }
        // `values` is the jagged concatenation of ids, so its length must equal sum(lengths).
        let mut total_ids = 0usize;
        for &len in &lengths {
            if len < 0 {
                return Err(NarsilError::InvalidRequest(
                    "DLRM sparse length must be non-negative".to_string(),
                ));
            }
            total_ids += len as usize;
        }
        if values.len() != total_ids {
            return Err(NarsilError::InvalidRequest(format!(
                "DLRM expected {total_ids} sparse values (sum of lengths), got {}",
                values.len()
            )));
        }

        Ok(PreparedTorchInput {
            dense_values,
            lengths,
            values,
            batch,
        })
    }

    /// Run one fused forward over `batch` samples (inputs already in key-major layout) and return
    /// the flat predictions.
    fn forward(
        &self,
        dense_values: &[f32],
        lengths_km: &[i64],
        values_km: &[i64],
        batch: usize,
    ) -> Result<Vec<f32>> {
        let dense = TchTensor::from_slice(dense_values)
            .reshape([batch as i64, self.inner.config.dense_in_features as i64])
            .to_device(self.inner.device);

        // KeyedJaggedTensor lengths/values are int32 in the TorchRec artifact.
        let lengths_i32 = to_i32(lengths_km)?;
        let values_i32 = to_i32(values_km)?;
        let lengths = TchTensor::from_slice(&lengths_i32).to_device(self.inner.device);
        let values = TchTensor::from_slice(&values_i32).to_device(self.inner.device);

        let batch = IValue::GenericDict(vec![
            (IValue::String(DENSE_KEY.to_string()), IValue::Tensor(dense)),
            (
                IValue::String(LENGTHS_KEY.to_string()),
                IValue::Tensor(lengths),
            ),
            (
                IValue::String(VALUES_KEY.to_string()),
                IValue::Tensor(values),
            ),
        ]);

        let output = self
            .inner
            .module
            .forward_is(&[batch])
            .map_err(|err| NarsilError::Backend(format!("torch DLRM forward failed: {err}")))?;

        let predictions = match output {
            IValue::GenericDict(pairs) => pairs
                .into_iter()
                .find_map(|(key, value)| match (key, value) {
                    (IValue::String(name), IValue::Tensor(tensor)) if name == OUTPUT_KEY => {
                        Some(tensor)
                    }
                    _ => None,
                })
                .ok_or_else(|| {
                    NarsilError::Backend(format!(
                        "torch DLRM output dict missing '{OUTPUT_KEY}' tensor"
                    ))
                })?,
            IValue::Tensor(tensor) => tensor,
            other => {
                return Err(NarsilError::Backend(format!(
                    "unexpected torch DLRM output: {other:?}"
                )));
            }
        };

        // Flatten to 1-D so a `[B]` or `[B, 1]` prediction tensor both decode cleanly.
        let predictions = predictions
            .to_kind(Kind::Float)
            .to_device(Device::Cpu)
            .reshape([-1]);
        Vec::<f32>::try_from(&predictions).map_err(|err| {
            NarsilError::Backend(format!("torch DLRM output conversion failed: {err}"))
        })
    }

    fn response_from_output(&self, batch: usize, output: Vec<f32>) -> Result<InferenceResponse> {
        let mut metadata = HashMap::new();
        metadata.insert("backend".to_string(), self.inner.name.to_string());
        metadata.insert("runtime".to_string(), self.inner.runtime.to_string());
        metadata.insert("workload".to_string(), "dlrm".to_string());

        Ok(InferenceResponse {
            outputs: vec![Tensor::from_f32(OUTPUT_KEY, vec![batch], output)?],
            metadata,
        })
    }

    fn infer_tensor_set(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        let prepared = self.prepare_request(&request)?;
        let sparse_features = self.inner.config.sparse_features;
        let (dense, lengths_km, values_km, total) =
            assemble_key_major(&[&prepared], sparse_features);
        let output = self.forward(&dense, &lengths_km, &values_km, total)?;
        self.response_from_output(prepared.batch, output)
    }

    /// Coalesce many requests into one key-major forward pass, then split outputs back per request.
    fn infer_batch_optimised(
        &self,
        requests: Vec<InferenceRequest>,
    ) -> Vec<Result<InferenceResponse>> {
        if requests.is_empty() {
            return vec![];
        }

        let mut prepared: Vec<Option<PreparedTorchInput>> = Vec::with_capacity(requests.len());
        let mut errors: Vec<Option<NarsilError>> = Vec::with_capacity(requests.len());
        for request in &requests {
            match self.prepare_request(request) {
                Ok(p) => {
                    prepared.push(Some(p));
                    errors.push(None);
                }
                Err(err) => {
                    prepared.push(None);
                    errors.push(Some(err));
                }
            }
        }

        let ok: Vec<&PreparedTorchInput> = prepared.iter().filter_map(Option::as_ref).collect();
        if ok.is_empty() {
            return errors
                .into_iter()
                .map(|err| {
                    Err(err.unwrap_or_else(|| NarsilError::Backend("empty DLRM batch".to_string())))
                })
                .collect();
        }

        let sparse_features = self.inner.config.sparse_features;
        let (dense, lengths_km, values_km, total) = assemble_key_major(&ok, sparse_features);
        let raw_outputs = match self.forward(&dense, &lengths_km, &values_km, total) {
            Ok(outputs) => outputs,
            Err(err) => {
                return errors
                    .into_iter()
                    .map(|maybe_err| {
                        Err(maybe_err.unwrap_or_else(|| {
                            NarsilError::Backend(format!(
                                "batched torch DLRM forward failed: {err}"
                            ))
                        }))
                    })
                    .collect();
            }
        };

        let mut offset = 0usize;
        prepared
            .into_iter()
            .zip(errors)
            .map(|(prepared, maybe_err)| {
                if let Some(err) = maybe_err {
                    return Err(err);
                }
                let batch = prepared.map(|p| p.batch).unwrap_or(0);
                let end = offset.checked_add(batch).ok_or_else(|| {
                    NarsilError::Backend("torch DLRM batch output offset overflow".to_string())
                })?;
                let output = raw_outputs.get(offset..end).ok_or_else(|| {
                    NarsilError::Backend("torch DLRM batch output slice out of bounds".to_string())
                })?;
                offset = end;
                self.response_from_output(batch, output.to_vec())
            })
            .collect()
    }
}

impl InferenceBackend for TorchDlrmBackend {
    fn metadata(&self) -> BackendMetadata {
        BackendMetadata {
            name: self.inner.name.to_string(),
            kind: self.inner.runtime.to_string(),
        }
    }

    fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        self.infer_tensor_set(request)
    }

    fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        self.infer_batch_optimised(requests)
    }
}

#[cfg(test)]
mod tests {
    use super::{PreparedTorchInput, assemble_key_major};

    fn prepared(lengths: Vec<i64>, values: Vec<i64>, batch: usize) -> PreparedTorchInput {
        PreparedTorchInput {
            dense_values: vec![0.0; batch],
            lengths,
            values,
            batch,
        }
    }

    #[test]
    fn key_major_single_request_one_id_per_feature() {
        // F=2, B=2, batch-major values [s0f0, s0f1, s1f0, s1f1].
        let p = prepared(vec![1, 1, 1, 1], vec![10, 11, 20, 21], 2);
        let (_dense, lengths_km, values_km, total) = assemble_key_major(&[&p], 2);
        assert_eq!(total, 2);
        // key-major: feature0 [s0, s1] = [10, 20], feature1 [s0, s1] = [11, 21].
        assert_eq!(values_km, vec![10, 20, 11, 21]);
        assert_eq!(lengths_km, vec![1, 1, 1, 1]);
    }

    #[test]
    fn key_major_coalesced_requests_interleave_by_feature() {
        // Two requests, F=2. reqA B=2, reqB B=1; total batch 3.
        let a = prepared(vec![1, 1, 1, 1], vec![100, 101, 110, 111], 2);
        let b = prepared(vec![1, 1], vec![200, 201], 1);
        let (_dense, _lengths_km, values_km, total) = assemble_key_major(&[&a, &b], 2);
        assert_eq!(total, 3);
        // feature0: a.s0, a.s1, b.s0 = 100, 110, 200; feature1: 101, 111, 201.
        assert_eq!(values_km, vec![100, 110, 200, 101, 111, 201]);
    }

    #[test]
    fn key_major_handles_jagged_lengths() {
        // F=2, B=1, jagged: feature0 has 2 ids, feature1 has 1.
        let p = prepared(vec![2, 1], vec![7, 8, 9], 1);
        let (_dense, lengths_km, values_km, _total) = assemble_key_major(&[&p], 2);
        assert_eq!(lengths_km, vec![2, 1]);
        assert_eq!(values_km, vec![7, 8, 9]);
    }
}
