use std::collections::HashMap;

use burn_flex::Flex;
use burn_tensor::{Tensor, TensorData};

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

type BurnCpu = Flex;

#[derive(Debug, Clone)]
pub struct BurnBackend {
    scale: f32,
    bias: f32,
}

impl BurnBackend {
    pub fn new() -> Self {
        Self {
            scale: 1.0,
            bias: 0.0,
        }
    }

    pub fn with_affine(scale: f32, bias: f32) -> Self {
        Self { scale, bias }
    }

    fn infer_tensor(&self, input: &crate::schema::Tensor) -> Result<crate::schema::Tensor> {
        let values = input.f32_values()?;
        let output_values = match input.shape.len() {
            1 => self.run_rank(values, [input.shape[0]])?,
            2 => self.run_rank(values, [input.shape[0], input.shape[1]])?,
            3 => self.run_rank(values, [input.shape[0], input.shape[1], input.shape[2]])?,
            4 => self.run_rank(
                values,
                [
                    input.shape[0],
                    input.shape[1],
                    input.shape[2],
                    input.shape[3],
                ],
            )?,
            rank => {
                return Err(NarsilError::Backend(format!(
                    "burn backend currently supports rank 1-4 tensors, got rank {rank}"
                )));
            }
        };

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

    fn run_rank<const D: usize>(&self, values: Vec<f32>, shape: [usize; D]) -> Result<Vec<f32>> {
        let device = Default::default();
        let data = TensorData::new(values, shape);
        let tensor = Tensor::<BurnCpu, D>::from_data(data, &device);
        let output = tensor * self.scale + self.bias;
        output
            .into_data()
            .to_vec::<f32>()
            .map_err(|err| NarsilError::Backend(format!("burn tensor conversion failed: {err}")))
    }
}

impl Default for BurnBackend {
    fn default() -> Self {
        Self::new()
    }
}

impl InferenceBackend for BurnBackend {
    fn metadata(&self) -> BackendMetadata {
        BackendMetadata {
            name: "burn-flex".to_string(),
            kind: "burn".to_string(),
        }
    }

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

        let output = self.infer_tensor(first)?;
        let metadata = HashMap::from([
            ("backend".to_string(), "burn".to_string()),
            ("runtime".to_string(), "burn-flex".to_string()),
        ]);

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