use std::collections::HashMap;

use burn_flex::{Flex, FlexDevice};
use burn_tensor::activation;
use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, TensorData};

use crate::backend::{BackendMetadata, InferenceBackend};
use crate::error::{NarsilError, Result};
use crate::schema::{InferenceRequest, InferenceResponse};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MlpConfig {
    pub input_dim: usize,
    pub hidden_dim: usize,
    pub output_dim: usize,
}

impl MlpConfig {
    pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Result<Self> {
        let config = Self {
            input_dim,
            hidden_dim,
            output_dim,
        };
        config.validate()?;
        Ok(config)
    }

    pub fn validate(&self) -> Result<()> {
        if self.input_dim == 0 || self.hidden_dim == 0 || self.output_dim == 0 {
            return Err(NarsilError::InvalidRequest(
                "MLP dimensions must be positive".to_string(),
            ));
        }
        Ok(())
    }
}

impl Default for MlpConfig {
    fn default() -> Self {
        Self {
            input_dim: 1024,
            hidden_dim: 2048,
            output_dim: 1000,
        }
    }
}

#[derive(Debug, Clone)]
pub struct BurnMlpModel<B: Backend<FloatElem = f32>> {
    config: MlpConfig,
    device: B::Device,
    w1: Tensor<B, 2>,
    w2: Tensor<B, 2>,
}

impl<B: Backend<FloatElem = f32>> BurnMlpModel<B> {
    pub fn new(config: MlpConfig, device: B::Device) -> Result<Self> {
        config.validate()?;

        let w1_values = deterministic_weights(config.input_dim, config.hidden_dim);
        let w2_values = deterministic_weights(config.hidden_dim, config.output_dim);
        let w1 = Tensor::<B, 2>::from_data(
            TensorData::new(w1_values, [config.input_dim, config.hidden_dim]),
            &device,
        );
        let w2 = Tensor::<B, 2>::from_data(
            TensorData::new(w2_values, [config.hidden_dim, config.output_dim]),
            &device,
        );

        Ok(Self {
            config,
            device,
            w1,
            w2,
        })
    }

    pub fn config(&self) -> MlpConfig {
        self.config
    }

    /// Load weights from a raw binary file (little-endian f32, row-major).
    /// Layout: w1 [input_dim × hidden_dim], then w2 [hidden_dim × output_dim].
    pub fn load_weights(
        path: &std::path::Path,
        config: &MlpConfig,
        device: &B::Device,
    ) -> crate::error::Result<Self> {
        let w1_count = config
            .input_dim
            .checked_mul(config.hidden_dim)
            .ok_or_else(|| {
                crate::error::NarsilError::Backend("MLP dimensions overflow".to_string())
            })?;
        let w2_count = config
            .hidden_dim
            .checked_mul(config.output_dim)
            .ok_or_else(|| {
                crate::error::NarsilError::Backend("MLP dimensions overflow".to_string())
            })?;
        let expected_bytes = w1_count
            .checked_add(w2_count)
            .and_then(|total| total.checked_mul(4))
            .ok_or_else(|| {
                crate::error::NarsilError::Backend("MLP dimensions overflow".to_string())
            })?;

        let bytes = std::fs::read(path).map_err(|err| {
            crate::error::NarsilError::Backend(format!(
                "failed to read model weights from {}: {err}",
                path.display()
            ))
        })?;

        if bytes.len() != expected_bytes {
            return Err(crate::error::NarsilError::Backend(format!(
                "model file size mismatch: expected {expected_bytes} bytes \
                 ({w1_count} + {w2_count} f32s), got {}",
                bytes.len()
            )));
        }

        let mut floats = bytes
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes(b.try_into().unwrap()));

        let w1_values: Vec<f32> = floats.by_ref().take(w1_count).collect();
        let w2_values: Vec<f32> = floats.collect();

        let w1 = Tensor::<B, 2>::from_data(
            TensorData::new(w1_values, [config.input_dim, config.hidden_dim]),
            device,
        );
        let w2 = Tensor::<B, 2>::from_data(
            TensorData::new(w2_values, [config.hidden_dim, config.output_dim]),
            device,
        );

        Ok(Self {
            config: *config,
            device: device.clone(),
            w1,
            w2,
        })
    }

    pub fn infer_batch_f32(&self, values: Vec<f32>, batch_size: usize) -> Result<Vec<f32>> {
        if batch_size == 0 {
            return Err(NarsilError::InvalidRequest(
                "MLP batch size must be positive".to_string(),
            ));
        }

        let expected_values = batch_size
            .checked_mul(self.config.input_dim)
            .ok_or_else(|| NarsilError::InvalidRequest("MLP input size overflow".to_string()))?;
        if values.len() != expected_values {
            return Err(NarsilError::InvalidRequest(format!(
                "MLP expected {expected_values} input values for batch {batch_size} and input dim {}, got {}",
                self.config.input_dim,
                values.len()
            )));
        }

        let input = Tensor::<B, 2>::from_data(
            TensorData::new(values, [batch_size, self.config.input_dim]),
            &self.device,
        );
        let output = self.forward(input);
        output.into_data().to_vec::<f32>().map_err(|err| {
            NarsilError::Backend(format!("burn MLP output conversion failed: {err}"))
        })
    }

    pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        let hidden = activation::relu(input.matmul(self.w1.clone()) + 0.01);
        hidden.matmul(self.w2.clone()) + 0.02
    }
}

#[derive(Debug, Clone)]
pub struct BurnMlpBackend<B: Backend<FloatElem = f32>> {
    model: BurnMlpModel<B>,
    name: &'static str,
    runtime: &'static str,
}

impl BurnMlpBackend<Flex> {
    pub fn flex(config: MlpConfig) -> Result<Self> {
        Self::new(config, FlexDevice, "burn-flex-mlp", "burn-flex")
    }

    /// Load weights from disk using the CPU (Flex) backend. Convenience wrapper around `load()`.
    pub fn flex_load(path: &std::path::Path, config: MlpConfig) -> crate::error::Result<Self> {
        Self::load(path, config, FlexDevice, "burn-flex-mlp", "burn-flex")
    }
}

impl<B: Backend<FloatElem = f32>> BurnMlpBackend<B> {
    pub fn new(
        config: MlpConfig,
        device: B::Device,
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        Ok(Self {
            model: BurnMlpModel::new(config, device)?,
            name,
            runtime,
        })
    }

    pub fn config(&self) -> MlpConfig {
        self.model.config()
    }

    /// Create a backend with weights loaded from a binary file.
    /// Use `flex()` or `new()` if you want deterministic toy weights.
    pub fn load(
        path: &std::path::Path,
        config: MlpConfig,
        device: B::Device,
        name: &'static str,
        runtime: &'static str,
    ) -> crate::error::Result<Self> {
        Ok(Self {
            model: BurnMlpModel::load_weights(path, &config, &device)?,
            name,
            runtime,
        })
    }

    fn infer_tensor(&self, input: &crate::schema::Tensor) -> Result<crate::schema::Tensor> {
        let values = input.f32_values()?;
        let batch_size = match input.shape.as_slice() {
            [input_dim] if *input_dim == self.config().input_dim => 1,
            [batch_size, input_dim] if *input_dim == self.config().input_dim => *batch_size,
            _ => {
                return Err(NarsilError::InvalidRequest(format!(
                    "MLP expects input shape [{}] or [batch, {}], got {:?}",
                    self.config().input_dim,
                    self.config().input_dim,
                    input.shape
                )));
            }
        };

        let output_values = self.model.infer_batch_f32(values, batch_size)?;
        let output_shape = if input.shape.len() == 1 {
            vec![self.config().output_dim]
        } else {
            vec![batch_size, self.config().output_dim]
        };

        crate::schema::Tensor::from_f32(
            if input.name.is_empty() {
                "output".to_string()
            } else {
                format!("{}_mlp_output", input.name)
            },
            output_shape,
            output_values,
        )
    }
}

impl<B: Backend<FloatElem = f32>> BurnMlpBackend<B> {
    fn infer_batch_optimised(
        &self,
        requests: Vec<crate::schema::InferenceRequest>,
    ) -> Vec<crate::error::Result<crate::schema::InferenceResponse>> {
        if requests.is_empty() {
            return vec![];
        }

        // Collect and validate all inputs.
        let mut all_inputs: Vec<f32> = Vec::new();
        let mut batch_sizes: Vec<usize> = Vec::new();
        let mut names: Vec<String> = Vec::new();
        let mut errors: Vec<Option<crate::error::NarsilError>> = Vec::new();

        let input_dim = self.config().input_dim;
        for req in &requests {
            if req.inputs.len() != 1 {
                errors.push(Some(crate::error::NarsilError::Backend(
                    "batch inference requires exactly 1 input tensor per request".to_string(),
                )));
                batch_sizes.push(0);
                names.push(String::new());
                continue;
            }
            let tensor = &req.inputs[0];
            let batch = match tensor.shape.as_slice() {
                [dim] if *dim == input_dim => 1,
                [b, dim] if *dim == input_dim => *b,
                _ => {
                    errors.push(Some(crate::error::NarsilError::InvalidRequest(format!(
                        "MLP expects input shape [{}] or [batch, {}], got {:?}",
                        input_dim, input_dim, tensor.shape
                    ))));
                    batch_sizes.push(0);
                    names.push(String::new());
                    continue;
                }
            };
            match tensor.f32_values() {
                Ok(values) => {
                    all_inputs.extend_from_slice(&values);
                    batch_sizes.push(batch);
                    let output_name = if tensor.name.is_empty() {
                        "output".to_string()
                    } else {
                        format!("{}_mlp_output", tensor.name)
                    };
                    names.push(output_name);
                    errors.push(None);
                }
                Err(e) => {
                    errors.push(Some(e));
                    batch_sizes.push(0);
                    names.push(String::new());
                }
            }
        }

        let total_batch: usize = batch_sizes.iter().sum();
        if total_batch == 0 {
            return errors
                .into_iter()
                .map(|e| {
                    Err(e.unwrap_or_else(|| {
                        crate::error::NarsilError::Backend("empty batch".to_string())
                    }))
                })
                .collect();
        }

        // Single batched forward pass.
        let output_dim = self.model.config.output_dim;
        let raw_outputs = match self.model.infer_batch_f32(all_inputs, total_batch) {
            Ok(o) => o,
            Err(e) => {
                return errors
                    .into_iter()
                    .map(|_| Err(crate::error::NarsilError::Backend(e.to_string())))
                    .collect();
            }
        };

        // Split outputs and wrap in responses.
        let mut offset = 0usize;
        errors
            .into_iter()
            .zip(batch_sizes.iter())
            .zip(names.iter())
            .map(|((err, &batch), name)| {
                if let Some(e) = err {
                    return Err(e);
                }
                let count = batch * output_dim;
                let slice = raw_outputs.get(offset..offset + count).ok_or_else(|| {
                    crate::error::NarsilError::Backend(
                        "batch output slice out of bounds (accounting bug)".to_string(),
                    )
                })?;
                offset += count;
                let shape = if batch == 1 {
                    vec![output_dim]
                } else {
                    vec![batch, output_dim]
                };
                let out_tensor = crate::schema::Tensor::from_f32(name, shape, slice.to_vec())?;
                Ok(crate::schema::InferenceResponse {
                    outputs: vec![out_tensor],
                    metadata: std::collections::HashMap::new(),
                })
            })
            .collect()
    }
}

impl<B: Backend<FloatElem = f32>> InferenceBackend for BurnMlpBackend<B> {
    fn metadata(&self) -> BackendMetadata {
        BackendMetadata {
            name: self.name.to_string(),
            kind: "burn-mlp".to_string(),
        }
    }

    fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        if request.inputs.len() != 1 {
            return Err(NarsilError::InvalidRequest(format!(
                "MLP backend requires exactly one input tensor, got {}",
                request.inputs.len()
            )));
        }
        let first = &request.inputs[0];

        let output = self.infer_tensor(first)?;
        let config = self.config();
        let metadata = HashMap::from([
            ("backend".to_string(), "burn".to_string()),
            ("runtime".to_string(), self.runtime.to_string()),
            ("workload".to_string(), "mlp".to_string()),
            ("input_dim".to_string(), config.input_dim.to_string()),
            ("hidden_dim".to_string(), config.hidden_dim.to_string()),
            ("output_dim".to_string(), config.output_dim.to_string()),
        ]);

        Ok(InferenceResponse {
            outputs: vec![output],
            metadata,
        })
    }

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

pub fn deterministic_input(batch_size: usize, input_dim: usize) -> Vec<f32> {
    (0..batch_size * input_dim)
        .map(|index| ((index % 251) as f32 / 251.0) - 0.5)
        .collect()
}

pub fn deterministic_weights(rows: usize, cols: usize) -> Vec<f32> {
    let scale = (rows as f32).sqrt().recip();
    (0..rows * cols)
        .map(|index| {
            let value = ((index * 37 + 17) % 1009) as f32 / 1009.0;
            (value - 0.5) * scale
        })
        .collect()
}
