use crate::errors::*;
use crate::faceembed::facenet::FaceNetEmbedder;
use crate::ort_ep::*;
use error_stack::ResultExt;
use ndarray::{Array2, ArrayView4};
use ort::{
    execution_providers::{CPUExecutionProvider, ExecutionProviderDispatch},
    session::{Session, builder::GraphOptimizationLevel},
    value::Tensor,
};
use std::path::Path;

#[derive(Debug)]
pub struct EmbeddingGenerator {
    session: Session,
}

pub struct EmbeddingGeneratorBuilder {
    model_data: Vec<u8>,
    execution_providers: Option<Vec<ExecutionProviderDispatch>>,
    intra_threads: Option<usize>,
    inter_threads: Option<usize>,
}

impl EmbeddingGeneratorBuilder {
    pub fn new(model: impl AsRef<[u8]>) -> crate::errors::Result<Self> {
        Ok(Self {
            model_data: model.as_ref().to_vec(),
            execution_providers: None,
            intra_threads: None,
            inter_threads: None,
        })
    }

    pub fn with_execution_providers(mut self, providers: impl AsRef<[ExecutionProvider]>) -> Self {
        let execution_providers: Vec<ExecutionProviderDispatch> = providers
            .as_ref()
            .iter()
            .filter_map(|provider| provider.to_dispatch())
            .collect();

        if !execution_providers.is_empty() {
            self.execution_providers = Some(execution_providers);
        } else {
            tracing::warn!("No valid execution providers found, falling back to CPU");
            self.execution_providers = Some(vec![CPUExecutionProvider::default().build()]);
        }
        self
    }

    pub fn with_intra_threads(mut self, threads: usize) -> Self {
        self.intra_threads = Some(threads);
        self
    }

    pub fn with_inter_threads(mut self, threads: usize) -> Self {
        self.inter_threads = Some(threads);
        self
    }

    pub fn build(self) -> crate::errors::Result<EmbeddingGenerator> {
        let mut session_builder = Session::builder()
            .change_context(Error)
            .attach_printable("Failed to create session builder")?;

        // Set execution providers
        if let Some(providers) = self.execution_providers {
            session_builder = session_builder
                .with_execution_providers(providers)
                .change_context(Error)
                .attach_printable("Failed to set execution providers")?;
        } else {
            // Default to CPU
            session_builder = session_builder
                .with_execution_providers([CPUExecutionProvider::default().build()])
                .change_context(Error)
                .attach_printable("Failed to set default CPU execution provider")?;
        }

        // Set threading options
        if let Some(threads) = self.intra_threads {
            session_builder = session_builder
                .with_intra_threads(threads)
                .change_context(Error)
                .attach_printable("Failed to set intra threads")?;
        }

        if let Some(threads) = self.inter_threads {
            session_builder = session_builder
                .with_inter_threads(threads)
                .change_context(Error)
                .attach_printable("Failed to set inter threads")?;
        }

        // Set optimization level
        session_builder = session_builder
            .with_optimization_level(GraphOptimizationLevel::Level3)
            .change_context(Error)
            .attach_printable("Failed to set optimization level")?;

        // Create session from model bytes
        let session = session_builder
            .commit_from_memory(&self.model_data)
            .change_context(Error)
            .attach_printable("Failed to create ORT session from model bytes")?;

        tracing::info!("Successfully created ORT FaceNet session");

        Ok(EmbeddingGenerator { session })
    }
}

impl EmbeddingGenerator {
    const INPUT_NAME: &'static str = "serving_default_input_6:0";
    const OUTPUT_NAME: &'static str = "StatefulPartitionedCall:0";

    pub fn builder<T: AsRef<[u8]>>(
        model: T,
    ) -> std::result::Result<EmbeddingGeneratorBuilder, error_stack::Report<crate::errors::Error>>
    {
        EmbeddingGeneratorBuilder::new(model)
    }

    pub fn new(path: impl AsRef<Path>) -> crate::errors::Result<Self> {
        let model = std::fs::read(path)
            .change_context(Error)
            .attach_printable("Failed to read model file")?;
        Self::new_from_bytes(&model)
    }

    pub fn new_from_bytes(model: impl AsRef<[u8]>) -> crate::errors::Result<Self> {
        tracing::info!("Loading ORT face embedding model from bytes");
        Self::builder(model)?.build()
    }

    pub fn run_models(&mut self, faces: ArrayView4<u8>) -> crate::errors::Result<Array2<f32>> {
        // Convert input from u8 to f32 and normalize to [0, 1] range
        let input_tensor = crate::faceembed::preprocessing::preprocess(faces);

        // face_array = np.asarray(face_resized, 'float32')
        //  mean, std = face_array.mean(), face_array.std()
        //  face_normalized = (face_array - mean) / std
        // let input_tensor = faces.mean()

        tracing::trace!("Input tensor shape: {:?}", input_tensor.shape());

        // Create ORT input tensor
        let input_value = Tensor::from_array(input_tensor)
            .change_context(Error)
            .attach_printable("Failed to create input tensor")?;

        // Run inference
        tracing::debug!("Running ORT FaceNet inference");
        let outputs = self
            .session
            .run(ort::inputs![Self::INPUT_NAME => input_value])
            .change_context(Error)
            .attach_printable("Failed to run inference")?;

        // Extract output tensor
        let output = outputs
            .get(Self::OUTPUT_NAME)
            .ok_or(Error)
            .attach_printable("Missing output from FaceNet model")?
            .try_extract_tensor::<f32>()
            .change_context(Error)
            .attach_printable("Failed to extract output tensor")?;

        let (output_shape, output_data) = output;

        tracing::trace!("Output shape: {:?}", output_shape);

        // Convert to ndarray format
        let output_dims = output_shape.as_ref();

        // FaceNet typically outputs embeddings as [batch_size, embedding_dim]
        let batch_size = output_dims[0] as usize;
        let embedding_dim = output_dims[1] as usize;

        let output_array =
            ndarray::Array2::from_shape_vec((batch_size, embedding_dim), output_data.to_vec())
                .change_context(Error)
                .attach_printable("Failed to create output ndarray")?;

        tracing::trace!(
            "Generated embeddings with shape: {:?}",
            output_array.shape()
        );

        Ok(output_array)
    }
}

impl FaceNetEmbedder for EmbeddingGenerator {
    fn run_model(&mut self, faces: ArrayView4<u8>) -> crate::errors::Result<Array2<f32>> {
        self.run_models(faces)
    }
}

// Main trait implementation for backward compatibility
impl crate::faceembed::FaceEmbedder for EmbeddingGenerator {
    fn run_models(&mut self, faces: ArrayView4<u8>) -> crate::errors::Result<Array2<f32>> {
        // Need to create a mutable reference for the session
        // This is a workaround for the trait signature mismatch
        self.run_models(faces)
    }
}
