use std::path::PathBuf;

use anyhow::{Context, Result};
use ndarray::ArrayView3;
use ort::session::Session;
use ort::value::TensorRef;
use rayon::prelude::*;

use crate::database::{self, image_pair_to_pair_id};
use crate::models::{self, Device};
use crate::ransac::{self, RansacResult};

pub struct ExhaustiveMatcherConfig {
  pub database_path: PathBuf,
  pub lightglue_model: PathBuf,
  pub min_score: f32,
  pub device: Device,
}

const MIN_INLIERS: usize = 15;
/// two_view_geometries.config: UNCALIBRATED (uses F, no known intrinsics).
const UNCALIBRATED_CONFIG: i64 = 3;

struct ImageFeatures {
  image_id: i64,
  name: String,
  /// Original pixel-space [x, y, scale, orientation], as stored in the database.
  keypoints: Vec<[f32; 4]>,
  /// Flat n*2 keypoints re-normalized to [-1, 1], matching ALIKED's own convention
  /// (normalized by max(width, height) of the source image, not per-axis).
  kpts_norm: Vec<f32>,
  /// Flat row-major n*desc_dim descriptors.
  descriptors: Vec<f32>,
  desc_dim: usize,
}

fn load_features(conn: &rusqlite::Connection, image_id: i64, name: String) -> Result<ImageFeatures> {
  let keypoints = database::read_keypoints(conn, image_id)
    .with_context(|| format!("reading keypoints for {name} (run feature_extractor first?)"))?;
  let (rows, desc_dim, descriptors) = database::read_descriptors(conn, image_id)?;
  if rows != keypoints.len() {
    anyhow::bail!(
      "{name}: keypoint count ({}) doesn't match descriptor row count ({rows})",
      keypoints.len()
    );
  }

  let (w, h) = database::image_dimensions(conn, image_id)?;
  let max_dim = w.max(h) as f32;
  let kpts_norm: Vec<f32> = keypoints
    .iter()
    .flat_map(|k| [2.0 * k[0] / max_dim - 1.0, 2.0 * k[1] / max_dim - 1.0])
    .collect();

  Ok(ImageFeatures {
    image_id,
    name,
    keypoints,
    kpts_norm,
    descriptors,
    desc_dim,
  })
}

/// Runs LightGlue on a pair of images' features. Returns (index_in_a, index_in_b) pairs
/// for matches passing `min_score`, with unmatched (-1) entries already excluded.
fn run_lightglue(
  session: &mut Session,
  a: &ImageFeatures,
  b: &ImageFeatures,
  min_score: f32,
) -> Result<Vec<(usize, usize)>> {
  let n = a.keypoints.len();
  let m = b.keypoints.len();
  if n == 0 || m == 0 {
    return Ok(Vec::new());
  }

  let kpts0 = ArrayView3::from_shape((1, n, 2), &a.kpts_norm)?;
  let kpts1 = ArrayView3::from_shape((1, m, 2), &b.kpts_norm)?;
  let desc0 = ArrayView3::from_shape((1, n, a.desc_dim), &a.descriptors)?;
  let desc1 = ArrayView3::from_shape((1, m, b.desc_dim), &b.descriptors)?;

  let outputs = session.run(ort::inputs! {
      "kpts0" => TensorRef::from_array_view(kpts0)?,
      "kpts1" => TensorRef::from_array_view(kpts1)?,
      "desc0" => TensorRef::from_array_view(desc0)?,
      "desc1" => TensorRef::from_array_view(desc1)?,
  })?;

  let matches_val = outputs
    .get("matches0")
    .ok_or_else(|| models::missing_output_err("matches0", &outputs))?;
  let (_shape, matches_data) = matches_val.try_extract_tensor::<i64>()?;

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

  Ok(
    matches_data
      .iter()
      .zip(scores_data.iter())
      .enumerate()
      .filter_map(|(i, (&m_idx, &score))| (m_idx >= 0 && score >= min_score).then_some((i, m_idx as usize)))
      .collect(),
  )
}

struct PendingPair {
  pair_id: i64,
  matched: Vec<(usize, usize)>,
  pts0: Vec<(f32, f32)>,
  pts1: Vec<(f32, f32)>,
}

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

  let images = database::all_images(&conn)?;
  if images.len() < 2 {
    anyhow::bail!(
      "need at least 2 images with extracted features to match; found {}. Run feature_extractor first.",
      images.len()
    );
  }

  let mut features = Vec::with_capacity(images.len());
  for (image_id, name) in images {
    features.push(load_features(&conn, image_id, name)?);
  }
  let n_images = features.len();

  println!(
    "Matching {n_images} images exhaustively ({} pairs)...",
    n_images * (n_images - 1) / 2
  );

  let mut pending: Vec<PendingPair> = Vec::new();
  for i in 0..n_images {
    for j in (i + 1)..n_images {
      let matched = run_lightglue(&mut session, &features[i], &features[j], cfg.min_score)?;
      if matched.is_empty() {
        continue;
      }

      // COLMAP's pair_id encoding requires the matches columns to be ordered
      // (smaller image_id, larger image_id); this normally coincides with insertion
      // order but is enforced explicitly here rather than assumed.
      let (lo, hi, swapped) = if features[i].image_id <= features[j].image_id {
        (&features[i], &features[j], false)
      } else {
        (&features[j], &features[i], true)
      };
      let ordered_matched: Vec<(usize, usize)> = if swapped {
        matched.iter().map(|&(a, b)| (b, a)).collect()
      } else {
        matched.clone()
      };

      let pair_id = image_pair_to_pair_id(lo.image_id, hi.image_id);
      let raw_matches: Vec<[i32; 2]> = ordered_matched.iter().map(|&(a, b)| [a as i32, b as i32]).collect();
      database::write_matches(&conn, pair_id, &raw_matches)?;

      let pts0: Vec<(f32, f32)> = ordered_matched
        .iter()
        .map(|&(a, _)| (lo.keypoints[a][0], lo.keypoints[a][1]))
        .collect();
      let pts1: Vec<(f32, f32)> = ordered_matched
        .iter()
        .map(|&(_, b)| (hi.keypoints[b][0], hi.keypoints[b][1]))
        .collect();

      println!(
        "  {} <-> {}: {} raw matches",
        features[i].name,
        features[j].name,
        ordered_matched.len()
      );

      pending.push(PendingPair {
        pair_id,
        matched: ordered_matched,
        pts0,
        pts1,
      });
    }
  }

  println!(
    "Running RANSAC geometric verification on {} candidate pairs...",
    pending.len()
  );
  let ransac_results: Vec<Option<RansacResult>> = pending
    .par_iter()
    .map(|p| {
      ransac::ransac_fundamental_matrix(
        &p.pts0,
        &p.pts1,
        ransac::DEFAULT_MAX_ITERATIONS,
        ransac::DEFAULT_INLIER_THRESHOLD,
        ransac::DEFAULT_CONFIDENCE,
      )
    })
    .collect();

  let mut verified_pairs = 0usize;
  let mut total_inliers = 0usize;
  for (pair, result) in pending.iter().zip(ransac_results.iter()) {
    let Some(result) = result else { continue };
    if result.inliers.len() < MIN_INLIERS {
      continue;
    }
    let inlier_matches: Vec<[i32; 2]> = result
      .inliers
      .iter()
      .map(|&idx| {
        let (a, b) = pair.matched[idx];
        [a as i32, b as i32]
      })
      .collect();
    database::write_two_view_geometry(
      &conn,
      pair.pair_id,
      UNCALIBRATED_CONFIG,
      &inlier_matches,
      &result.f_matrix,
    )?;
    verified_pairs += 1;
    total_inliers += inlier_matches.len();
  }

  let avg_inliers = if verified_pairs > 0 {
    total_inliers as f64 / verified_pairs as f64
  } else {
    0.0
  };
  println!(
    "Done: {} pairs had raw matches, {verified_pairs} passed geometric verification (avg {avg_inliers:.1} inliers/pair)",
    pending.len()
  );

  Ok(())
}