use anyhow::{Context, Result};
use image::DynamicImage;
use ndarray::Array4;
use ort::session::Session;
use ort::value::Tensor;
use std::path::PathBuf;
 
const MODEL_URL: &str =
    "https://huggingface.co/onnx-community/siglip2-base-patch16-224-ONNX/resolve/main/onnx/vision_model_q4.onnx";
const CACHE: &str = "vision_model_q4.onnx";
 
fn preprocess(img: &DynamicImage) -> Array4<f32> {
    let img = img.resize_exact(224, 224, image::imageops::FilterType::CatmullRom).to_rgb8();
    let (w, h) = img.dimensions();
    let mut tensor = Array4::zeros((1, 3, h as usize, w as usize));
 
    for (x, y, p) in img.enumerate_pixels() {
        tensor[[0, 0, y as usize, x as usize]] = (p[0] as f32 / 255.0 - 0.5) / 0.5;
        tensor[[0, 1, y as usize, x as usize]] = (p[1] as f32 / 255.0 - 0.5) / 0.5;
        tensor[[0, 2, y as usize, x as usize]] = (p[2] as f32 / 255.0 - 0.5) / 0.5;
    }
 
    tensor
}
 
async fn download_model() -> Result<PathBuf> {
    let path = PathBuf::from(CACHE);
    if path.exists() {
        return Ok(path);
    }
 
    eprintln!("downloading model from huggingface...");
    let resp = reqwest::get(MODEL_URL).await?;
    let bytes = resp.bytes().await?;
    std::fs::write(&path, &bytes)?;
    eprintln!("model cached to {}", path.display());
    Ok(path)
}
 
fn embed_image(session: &mut Session, img: &DynamicImage) -> Result<Vec<f32>> {
    let tensor = preprocess(img);
    let input_tensor = Tensor::from_array(tensor)?;
 
    let outputs = session.run(ort::inputs!["pixel_values" => input_tensor])?;
 
    if let Some(out) = outputs.get("pooler_output") {
        let arr = out.try_extract_array::<f32>()?;
        Ok(arr.iter().copied().collect())
    } else if let Some(out) = outputs.get("last_hidden_state") {
        let arr = out.try_extract_array::<f32>()?;
        let shape = arr.shape();
        if shape.len() == 3 && shape[0] == 1 {
            let dim = shape[2];
            Ok((0..dim).map(|i| arr[[0, 0, i]]).collect())
        } else {
            Ok(arr.iter().copied().collect())
        }
    } else {
        anyhow::bail!("no supported output tensor found");
    }
}
 
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum();
    let nb: f32 = b.iter().map(|x| x * x).sum();
    dot / (na.sqrt() * nb.sqrt())
}
 
#[tokio::main]
async fn main() -> Result<()> {
    let path_a = std::env::args().nth(1).unwrap_or_else(|| "a.jpg".into());
    let path_b = std::env::args().nth(2).unwrap_or_else(|| "b.jpg".into());
 
    let model_path = download_model().await?;
 
    let mut builder = Session::builder().map_err(|e| anyhow::anyhow!("{}", e))?;
    builder = builder.with_intra_threads(4).map_err(|e| anyhow::anyhow!("{}", e))?;
    let mut session = builder.commit_from_file(model_path)?;
 
    let img_a = image::open(&path_a).context("failed to open first image")?;
    let img_b = image::open(&path_b).context("failed to open second image")?;
 
    let emb_a = embed_image(&mut session, &img_a)?;
    let emb_b = embed_image(&mut session, &img_b)?;
 
    eprintln!("embedding dim: {}", emb_a.len());
    println!("{}", cosine_similarity(&emb_a, &emb_b));
 
    Ok(())
}