use narsil_engine::schema::proto::inference_client::InferenceClient;
use narsil_engine::schema::proto::{DataType, PredictRequest, Tensor as ProtoTensor};
use narsil_engine::schema::{
    DataType as NativeDataType, InferenceRequest as NativeInferenceRequest, Tensor as NativeTensor,
};
use narsil_engine::{
    BurnBackend, BurnDlrmBackend, BurnMlpBackend, DlrmConfig, GrpcInferenceService,
    InferenceBackend, MlpConfig, WorkerPoolConfig,
};
use tokio_stream::wrappers::TcpListenerStream;

async fn spawn_server(
    service: GrpcInferenceService<impl narsil_engine::InferenceBackend + Clone>,
) -> String {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(service.into_server())
            .serve_with_incoming(TcpListenerStream::new(listener))
            .await
            .unwrap();
    });

    format!("http://{addr}")
}

fn f32_bytes(values: &[f32]) -> Vec<u8> {
    values
        .iter()
        .flat_map(|value| value.to_le_bytes())
        .collect()
}

fn f32_values(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
        .collect()
}

fn i32_bytes(values: &[i32]) -> Vec<u8> {
    values
        .iter()
        .flat_map(|value| value.to_le_bytes())
        .collect()
}

#[tokio::test]
async fn burn_backend_serves_tonic_predict() {
    let endpoint = spawn_server(GrpcInferenceService::new(BurnBackend::with_affine(
        2.0, 1.0,
    )))
    .await;
    let mut client = InferenceClient::connect(endpoint).await.unwrap();

    let response = client
        .predict(PredictRequest {
            model_name: "burn-affine".to_string(),
            inputs: vec![ProtoTensor {
                name: "input".to_string(),
                dtype: DataType::F32 as i32,
                shape: vec![2, 2],
                data: f32_bytes(&[1.0, 2.0, 3.0, 4.0]),
            }],
            parameters: Default::default(),
        })
        .await
        .unwrap()
        .into_inner();

    assert_eq!(response.outputs.len(), 1);
    assert_eq!(response.outputs[0].shape, vec![2, 2]);
    assert_eq!(
        f32_values(&response.outputs[0].data),
        vec![3.0, 5.0, 7.0, 9.0]
    );
    assert_eq!(response.metadata.get("backend").unwrap(), "burn");
}

#[tokio::test]
async fn burn_mlp_backend_serves_tonic_predict() {
    let config = MlpConfig::new(4, 8, 3).unwrap();
    let service = GrpcInferenceService::new(BurnMlpBackend::flex(config).unwrap());
    let endpoint = spawn_server(service).await;
    let mut client = InferenceClient::connect(endpoint).await.unwrap();

    let response = client
        .predict(PredictRequest {
            model_name: "burn-mlp".to_string(),
            inputs: vec![ProtoTensor {
                name: "input".to_string(),
                dtype: DataType::F32 as i32,
                shape: vec![4],
                data: f32_bytes(&[0.1, 0.2, 0.3, 0.4]),
            }],
            parameters: Default::default(),
        })
        .await
        .unwrap()
        .into_inner();

    assert_eq!(response.outputs.len(), 1);
    assert_eq!(response.outputs[0].shape, vec![3]);
    assert_eq!(f32_values(&response.outputs[0].data).len(), 3);
    assert_eq!(response.metadata.get("workload").unwrap(), "mlp");
}

#[tokio::test]
async fn burn_dlrm_backend_serves_tonic_predict() {
    let config = DlrmConfig {
        dense_in_features: 4,
        sparse_features: 3,
        num_embeddings: 16,
        embedding_dim: 4,
        dense_arch_layer_sizes: vec![8, 4],
        over_arch_layer_sizes: vec![8, 1],
    };
    let service = GrpcInferenceService::new(BurnDlrmBackend::flex(config).unwrap());
    let endpoint = spawn_server(service).await;
    let mut client = InferenceClient::connect(endpoint).await.unwrap();

    let response = client
        .predict(PredictRequest {
            model_name: "burn-dlrm".to_string(),
            inputs: vec![
                ProtoTensor {
                    name: "dense".to_string(),
                    dtype: DataType::F32 as i32,
                    shape: vec![2, 4],
                    data: f32_bytes(&[0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1]),
                },
                ProtoTensor {
                    name: "lengths".to_string(),
                    dtype: DataType::I32 as i32,
                    shape: vec![2, 3],
                    data: i32_bytes(&[1, 1, 1, 1, 1, 1]),
                },
                ProtoTensor {
                    name: "values".to_string(),
                    dtype: DataType::I32 as i32,
                    shape: vec![2, 3],
                    data: i32_bytes(&[1, 2, 3, 4, 5, 6]),
                },
            ],
            parameters: Default::default(),
        })
        .await
        .unwrap()
        .into_inner();

    assert_eq!(response.outputs.len(), 1);
    assert_eq!(response.outputs[0].shape, vec![2]);
    assert_eq!(f32_values(&response.outputs[0].data).len(), 2);
    assert_eq!(response.metadata.get("workload").unwrap(), "dlrm");
}

#[tokio::test]
async fn worker_pool_serves_tonic_predict() {
    let service = GrpcInferenceService::with_worker_pool(
        BurnBackend::with_affine(3.0, -1.0),
        WorkerPoolConfig::new(2, 8).unwrap(),
    )
    .unwrap();
    let endpoint = spawn_server(service).await;
    let mut client = InferenceClient::connect(endpoint).await.unwrap();

    let response = client
        .predict(PredictRequest {
            model_name: "burn-affine-worker".to_string(),
            inputs: vec![ProtoTensor {
                name: "input".to_string(),
                dtype: DataType::F32 as i32,
                shape: vec![2],
                data: f32_bytes(&[2.0, 4.0]),
            }],
            parameters: Default::default(),
        })
        .await
        .unwrap()
        .into_inner();

    assert_eq!(response.outputs.len(), 1);
    assert_eq!(f32_values(&response.outputs[0].data), vec![5.0, 11.0]);
}

fn native_tensor(values: &[f32]) -> NativeTensor {
    NativeTensor::new(
        "input",
        NativeDataType::F32,
        vec![values.len()],
        f32_bytes(values),
    )
    .unwrap()
}

fn native_request(inputs: Vec<NativeTensor>) -> NativeInferenceRequest {
    NativeInferenceRequest {
        model_name: "test".to_string(),
        inputs,
        parameters: Default::default(),
    }
}

fn native_dlrm_request(
    dense: &[f32],
    dense_shape: Vec<usize>,
    lengths: &[i32],
    values: &[i32],
    sparse_shape: Vec<usize>,
) -> NativeInferenceRequest {
    native_request(vec![
        NativeTensor::new("dense", NativeDataType::F32, dense_shape, f32_bytes(dense)).unwrap(),
        NativeTensor::new(
            "lengths",
            NativeDataType::I32,
            sparse_shape.clone(),
            i32_bytes(lengths),
        )
        .unwrap(),
        NativeTensor::new(
            "values",
            NativeDataType::I32,
            sparse_shape,
            i32_bytes(values),
        )
        .unwrap(),
    ])
}

fn first_output_values(response: &narsil_engine::schema::InferenceResponse) -> Vec<f32> {
    f32_values(&response.outputs[0].data)
}

#[test]
fn burn_dlrm_backend_batches_requests() {
    let config = DlrmConfig {
        dense_in_features: 4,
        sparse_features: 3,
        num_embeddings: 16,
        embedding_dim: 4,
        dense_arch_layer_sizes: vec![8, 4],
        over_arch_layer_sizes: vec![8, 1],
    };
    let backend = BurnDlrmBackend::flex(config).unwrap();
    let first = native_dlrm_request(
        &[0.1, 0.2, 0.3, 0.4],
        vec![1, 4],
        &[1, 1, 1],
        &[1, 2, 3],
        vec![1, 3],
    );
    let second = native_dlrm_request(
        &[0.4, 0.3, 0.2, 0.1, 0.2, 0.2, 0.2, 0.2],
        vec![2, 4],
        &[1, 1, 1, 1, 1, 1],
        &[4, 5, 6, 7, 8, 9],
        vec![2, 3],
    );

    let expected_first = backend.infer(first.clone()).unwrap();
    let expected_second = backend.infer(second.clone()).unwrap();
    let mut batched = backend.infer_batch(vec![first, second]).into_iter();
    let actual_first = batched.next().unwrap().unwrap();
    let actual_second = batched.next().unwrap().unwrap();
    assert!(batched.next().is_none());

    assert_eq!(actual_first.outputs[0].shape, vec![1]);
    assert_eq!(actual_second.outputs[0].shape, vec![2]);
    assert_eq!(
        first_output_values(&actual_first),
        first_output_values(&expected_first)
    );
    assert_eq!(
        first_output_values(&actual_second),
        first_output_values(&expected_second)
    );
}

#[test]
fn burn_backend_rejects_extra_inputs() {
    let backend = BurnBackend::default();
    let err = backend
        .infer(native_request(vec![
            native_tensor(&[1.0, 2.0]),
            native_tensor(&[3.0, 4.0]),
        ]))
        .unwrap_err();

    assert!(err.to_string().contains("exactly one input tensor"));
}

#[test]
fn burn_mlp_backend_rejects_extra_inputs() {
    let backend = BurnMlpBackend::flex(MlpConfig::new(2, 4, 2).unwrap()).unwrap();
    let err = backend
        .infer(native_request(vec![
            native_tensor(&[1.0, 2.0]),
            native_tensor(&[3.0, 4.0]),
        ]))
        .unwrap_err();

    assert!(err.to_string().contains("exactly one input tensor"));
}
