mod database;
mod exhaustive_matcher;
mod feature_extractor;
mod models;
mod ransac;

use std::path::PathBuf;

use anyhow::Result;
use clap::{Parser, Subcommand};

use models::Device;

/// COLMAP-compatible feature extraction and matching using ALIKED + LightGlue.
#[derive(Parser)]
#[command(name = "aliked-colmap")]
struct Cli {
  #[command(subcommand)]
  command: Command,
}

#[derive(Subcommand)]
enum Command {
  /// Extract ALIKED keypoints/descriptors for all images and populate a COLMAP database.
  #[command(name = "feature_extractor")]
  FeatureExtractor {
    #[arg(long)]
    database_path: PathBuf,
    #[arg(long)]
    image_path: PathBuf,
    #[arg(long, default_value = "./models/aliked-n16rot-top2k-640.onnx")]
    aliked_model: PathBuf,
    #[arg(long, default_value_t = 2048)]
    top_k: usize,
    #[arg(long, default_value_t = 640)]
    image_size: u32,
    #[arg(long, value_enum, default_value_t = Device::Gpu)]
    device: Device,
  },
  /// Exhaustively match all image pairs with LightGlue and geometrically verify them.
  #[command(name = "exhaustive_matcher")]
  ExhaustiveMatcher {
    #[arg(long)]
    database_path: PathBuf,
    #[arg(long, default_value = "./models/lightglue_for_aliked.onnx")]
    lightglue_model: PathBuf,
    #[arg(long, default_value_t = 0.0)]
    min_score: f32,
    #[arg(long, value_enum, default_value_t = Device::Gpu)]
    device: Device,
  },
}

fn main() -> Result<()> {
  let cli = Cli::parse();

  match cli.command {
    Command::FeatureExtractor {
      database_path,
      image_path,
      aliked_model,
      top_k,
      image_size,
      device,
    } => feature_extractor::run(feature_extractor::FeatureExtractorConfig {
      database_path,
      image_path,
      aliked_model,
      top_k,
      image_size,
      device,
    }),
    Command::ExhaustiveMatcher {
      database_path,
      lightglue_model,
      min_score,
      device,
    } => exhaustive_matcher::run(exhaustive_matcher::ExhaustiveMatcherConfig {
      database_path,
      lightglue_model,
      min_score,
      device,
    }),
  }
}