use anyhow::{Context, Result, anyhow};
use itertools::Itertools;
use log::info;
use memmap2::Mmap;
use numpy::ndarray::Array3;
use safetensors::serialize_to_file;
use safetensors::tensor::{Dtype, TensorView};
use search_rdf::data::item::{FieldRef, load_image_ndarray_from_bytes};
use search_rdf::model::Embed;
use search_rdf::model::multimodal::MultiModalInput;
use std::collections::HashMap;
use std::fs::{File, create_dir_all, remove_file};
use std::io::{BufWriter, Write};
use std::path::Path;

use search_rdf::data::Data;
use search_rdf::data::DataSource;
use search_rdf::model::{EmbeddingModel, EmbeddingParams};
use search_rdf::utils::progress_bar;

use crate::search_rdf::config::Config;
use crate::search_rdf::model::get_or_load_model_and_params;

pub fn run(config_path: &Path, force: bool) -> Result<()> {
    let config = Config::load(config_path)?;
    let config_dir = config_path
        .parent()
        .expect("Failed to get config directory");

    info!("Building embeddings in {}", config_dir.display());

    let Some(embed) = &config.embeddings else {
        info!("No embedding configuration found.");
        return Ok(());
    };

    // Store models
    let mut models = HashMap::new();

    info!("Generating embeddings for {} datasets...", embed.len());

    for embed_config in embed {
        if config_dir.join(&embed_config.output).exists() && !force {
            info!(
                "[SKIP] {} (output exists, use --force to rebuild)",
                embed_config.name
            );
            continue;
        }

        let (model, params) =
            get_or_load_model_and_params(&embed_config.model, &mut models, &config)?;

        info!("[BUILD] {}...", embed_config.name);

        build_embeddings(
            config_dir,
            &embed_config.data,
            &embed_config.output,
            embed_config.batch_size,
            params,
            model,
        )?;

        info!(
            "[OK] {} -> {}",
            embed_config.name,
            embed_config.output.display()
        );
    }

    Ok(())
}

fn get_text_fields<'f>(fields: &[FieldRef<'f>]) -> Result<Vec<&'f str>> {
    fields
        .iter()
        .map(|f| {
            if f.is_text() {
                Ok(f.as_str())
            } else {
                Err(anyhow!(
                    "Expected text field for text embedding model, but found {:?}",
                    f
                ))
            }
        })
        .collect()
}

fn get_image_fields(fields: &[FieldRef]) -> Result<Vec<Array3<u8>>> {
    fields
        .iter()
        .map(|f| {
            if f.is_image() {
                // Load image data (handles both Image and ImageInline)
                let bytes = f.load_data().context("Failed to load image data")?;
                load_image_ndarray_from_bytes(&bytes)
                    .context("Failed to convert image bytes to array")
            } else {
                Err(anyhow!(
                    "Expected image field for image embedding model, but found text field"
                ))
            }
        })
        .collect()
}

fn get_multimodal_fields(fields: &[FieldRef]) -> Result<Vec<MultiModalInput>> {
    fields
        .iter()
        .map(|f| {
            if f.is_text() {
                Ok(MultiModalInput::Text(f.as_str().to_string()))
            } else if f.is_image() {
                let bytes = f.load_data().context("Failed to load image data")?;
                let arr = load_image_ndarray_from_bytes(&bytes)
                    .context("Failed to convert image bytes to array")?;
                Ok(MultiModalInput::Image(arr))
            } else {
                Err(anyhow!(
                    "Unsupported field type for OpenCLIP model: {:?}",
                    f
                ))
            }
        })
        .collect()
}

fn build_embeddings(
    base_dir: &Path,
    dataset: &Path,
    output: &Path,
    batch_size: usize,
    params: &EmbeddingParams,
    model: &EmbeddingModel,
) -> Result<()> {
    // Load data
    let data = Data::load(&base_dir.join(dataset))
        .context(format!("Failed to load data from: {}", dataset.display()))?;

    info!("Loaded {} items", data.len());
    info!(
        "Embedding {} fields with batch size {}...",
        data.total_fields(),
        batch_size
    );

    // Create output directory
    if let Some(parent) = base_dir.join(output).parent() {
        create_dir_all(parent)?;
    }

    let temp_file_path = base_dir.join(output).with_added_extension("tmp");
    let (temp_file, skip) = if temp_file_path.exists() {
        let temp_bytes = unsafe { Mmap::map(&File::open(&temp_file_path)?)? };
        if !temp_bytes
            .len()
            .is_multiple_of(size_of::<f32>() * model.num_dimensions())
        {
            return Err(anyhow!(
                "Temporary file at {} has invalid size, please remove it and try again",
                temp_file_path.display()
            ));
        }
        info!(
            "Resuming from existing temporary file: {}",
            temp_file_path.display()
        );
        let skip = temp_bytes.len() / (size_of::<f32>() * model.num_dimensions());
        info!("Skipping {} already embedded fields", skip);
        let file = File::options().append(true).open(&temp_file_path)?;
        (file, skip)
    } else {
        info!(
            "Using temporary file for embeddings: {}",
            temp_file_path.display()
        );
        (File::create(&temp_file_path)?, 0)
    };
    let mut temp_file = BufWriter::new(temp_file);

    let pb = progress_bar(
        "Generating embeddings",
        Some(data.total_fields() as u64 - skip as u64),
    )?;

    for chunk in &data
        .items()
        .flat_map(|(_, fields)| fields)
        .skip(skip)
        .chunks(batch_size)
    {
        let chunk: Vec<_> = chunk.collect();

        // Embed based on model type and field types
        let embeddings = match model {
            EmbeddingModel::SentenceTransformer(m) => {
                let texts = get_text_fields(&chunk)?;
                m.embed(&texts, params)?
            }
            EmbeddingModel::Vllm(m) => {
                let texts = get_text_fields(&chunk)?;
                m.embed(&texts, params)?
            }
            EmbeddingModel::HuggingFaceImage(m) => {
                let images = get_image_fields(&chunk)?;
                m.embed(&images, params)?
            }
            EmbeddingModel::OpenClip(m) => {
                let inputs = get_multimodal_fields(&chunk)?;
                m.embed(&inputs, params)?
            }
        };

        pb.inc(embeddings.len() as u64);
        embeddings
            .into_iter()
            .flatten()
            .map(|f| f.to_le_bytes())
            .try_for_each(|b| temp_file.write_all(&b))
            .context("Failed to write embeddings to temporary file")?;
    }

    temp_file.flush()?;
    pb.finish_with_message("Embeddings generated");

    let temp_bytes = unsafe { Mmap::map(&File::open(&temp_file_path)?)? };

    // Create tensors
    let embedding_tensor = TensorView::new(
        Dtype::F32,
        vec![data.total_fields() as usize, model.num_dimensions()],
        &temp_bytes,
    )?;

    // Serialize with model metadata
    let modality = model
        .modality()
        .iter()
        .map(serde_plain::to_string)
        .collect::<std::result::Result<Vec<_>, _>>()?
        .join(",");
    serialize_to_file(
        [("embedding", embedding_tensor)],
        Some(HashMap::from([
            (String::from("model"), model.model_name().to_string()),
            (String::from("provider"), model.provider().to_string()),
            (String::from("modality"), modality),
        ])),
        &base_dir.join(output),
    )?;

    info!("Saved embeddings to {}", output.display());

    // Remove temporary file
    remove_file(&temp_file_path)?;
    info!(
        "Removed temporary embedding file at {}",
        temp_file_path.display()
    );

    Ok(())
}
