use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use image::imageops::FilterType;
use ndarray::Array4;
use ort::session::Session;
use ort::value::TensorRef;
use walkdir::WalkDir;

use crate::database;
use crate::models::{self, Device};

pub struct FeatureExtractorConfig {
  pub database_path: PathBuf,
  pub image_path: PathBuf,
  pub aliked_model: PathBuf,
  pub top_k: usize,
  pub image_size: u32,
  pub device: Device,
}

struct Preprocessed {
  tensor: Array4<f32>,
  orig_width: u32,
  orig_height: u32,
}

const IMAGE_EXTENSIONS: [&str; 3] = ["jpg", "jpeg", "png"];

fn collect_image_paths(root: &Path) -> Result<Vec<PathBuf>> {
  let mut paths: Vec<PathBuf> = WalkDir::new(root)
    .into_iter()
    .filter_map(|e| e.ok())
    .filter(|e| e.file_type().is_file())
    .filter(|e| {
      e.path()
        .extension()
        .and_then(|s| s.to_str())
        .map(|s| IMAGE_EXTENSIONS.iter().any(|ext| ext.eq_ignore_ascii_case(s)))
        .unwrap_or(false)
    })
    .map(|e| e.path().to_path_buf())
    .collect();
  // Sorted so image_id assignment order matches filename order, which colmap mapper expects.
  paths.sort();
  Ok(paths)
}

/// Letterbox-resizes (top-left aligned, matching the ALIKED-LightGlue-ONNX reference
/// implementation) into an `image_size` x `image_size` canvas and produces an NCHW RGB
/// float32 tensor normalized to [0, 1].
fn preprocess(path: &Path, image_size: u32) -> Result<Preprocessed> {
  let img = image::open(path)
    .with_context(|| format!("opening image {}", path.display()))?
    .into_rgb8();
  let (w, h) = img.dimensions();
  if w == 0 || h == 0 {
    anyhow::bail!("image {} has zero dimension", path.display());
  }

  let resize_scale = (image_size as f32 / h as f32).min(image_size as f32 / w as f32);
  let w_new = ((w as f32 * resize_scale) as u32).clamp(1, image_size);
  let h_new = ((h as f32 * resize_scale) as u32).clamp(1, image_size);
  let resized = image::imageops::resize(&img, w_new, h_new, FilterType::Triangle);

  let size = image_size as usize;
  let mut tensor = Array4::<f32>::zeros((1, 3, size, size));
  for y in 0..h_new {
    for x in 0..w_new {
      let px = resized.get_pixel(x, y);
      for c in 0..3 {
        tensor[[0, c, y as usize, x as usize]] = px[c] as f32 / 255.0;
      }
    }
  }

  Ok(Preprocessed {
    tensor,
    orig_width: w,
    orig_height: h,
  })
}

/// (keypoints in normalized [-1,1] canvas coordinates, flat row-major descriptors [n x 128], per-keypoint scores)
type AlikedOutput = (Vec<[f32; 2]>, Vec<f32>, Vec<f32>);

/// Runs ALIKED on a preprocessed tensor.
fn run_aliked(session: &mut Session, tensor: &Array4<f32>) -> Result<AlikedOutput> {
  let input = TensorRef::from_array_view(tensor)?;
  let outputs = session.run(ort::inputs![input])?;

  let kpts_val = outputs
    .get("keypoints")
    .ok_or_else(|| models::missing_output_err("keypoints", &outputs))?;
  let (_kpts_shape, kpts_data) = kpts_val.try_extract_tensor::<f32>()?;
  let keypoints: Vec<[f32; 2]> = kpts_data.chunks_exact(2).map(|c| [c[0], c[1]]).collect();

  let desc_val = outputs
    .get("descriptors")
    .ok_or_else(|| models::missing_output_err("descriptors", &outputs))?;
  let (_desc_shape, desc_data) = desc_val.try_extract_tensor::<f32>()?;

  let scores_val = outputs
    .get("scores")
    .ok_or_else(|| models::missing_output_err("scores", &outputs))?;
  let (_scores_shape, scores_data) = scores_val.try_extract_tensor::<f32>()?;

  Ok((keypoints, desc_data.to_vec(), scores_data.to_vec()))
}

pub fn run(cfg: FeatureExtractorConfig) -> Result<()> {
  let conn = database::open(&cfg.database_path)?;
  let mut session = models::load_session(&cfg.aliked_model, cfg.device)?;

  let paths = collect_image_paths(&cfg.image_path)?;
  if paths.is_empty() {
    anyhow::bail!("no images found under {}", cfg.image_path.display());
  }
  println!("Found {} images under {}", paths.len(), cfg.image_path.display());

  for path in &paths {
    let rel_name = path
      .strip_prefix(&cfg.image_path)
      .unwrap_or(path)
      .to_string_lossy()
      .replace('\\', "/");

    let pre = preprocess(path, cfg.image_size)?;
    let camera_id = database::get_or_create_camera(&conn, pre.orig_width, pre.orig_height)?;
    let image_id = database::insert_image(&conn, &rel_name, camera_id)?;

    let (keypoints, descriptors_flat, scores) = run_aliked(&mut session, &pre.tensor)?;
    let n = keypoints.len();
    if descriptors_flat.len() % n.max(1) != 0 {
      anyhow::bail!(
        "{}: descriptor count ({}) doesn't divide evenly by keypoint count ({})",
        rel_name,
        descriptors_flat.len(),
        n
      );
    }
    let desc_dim = if n == 0 { 128 } else { descriptors_flat.len() / n };

    // The ONNX model already caps output count near its export-time top-k; this only
    // trims further if more keypoints than requested were returned.
    let keep: Vec<usize> = if n > cfg.top_k {
      let mut idx: Vec<usize> = (0..n).collect();
      idx.sort_unstable_by(|&a, &b| scores[b].partial_cmp(&scores[a]).unwrap());
      idx.truncate(cfg.top_k);
      idx.sort_unstable();
      idx
    } else {
      (0..n).collect()
    };

    let max_dim = pre.orig_width.max(pre.orig_height) as f32;
    let mut kpt_rows = Vec::with_capacity(keep.len());
    let mut desc_rows = Vec::with_capacity(keep.len() * desc_dim);
    for &i in &keep {
      let [x, y] = keypoints[i];
      // ALIKED's exported graph outputs keypoints normalized to [-1, 1] over the padded
      // square canvas; this maps directly back to original pixel coordinates.
      let x_orig = (x + 1.0) * max_dim / 2.0;
      let y_orig = (y + 1.0) * max_dim / 2.0;
      kpt_rows.push([x_orig, y_orig, 0.0, 0.0]);
      desc_rows.extend_from_slice(&descriptors_flat[i * desc_dim..(i + 1) * desc_dim]);
    }

    database::write_keypoints(&conn, image_id, &kpt_rows)?;
    database::write_descriptors(&conn, image_id, keep.len(), desc_dim, &desc_rows)?;

    println!("{rel_name}: {} keypoints", keep.len());
  }

  Ok(())
}