use std::path::Path;

use anyhow::{Context, Result};
use rusqlite::Connection;

/// COLMAP's SIMPLE_RADIAL camera model id.
pub const SIMPLE_RADIAL_MODEL_ID: i64 = 2;

/// COLMAP caps image ids below this value; the pair_id encoding relies on it.
pub const MAX_NUM_IMAGES: i64 = 2_147_483_647;

const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS cameras (
    camera_id         INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    model             INTEGER NOT NULL,
    width             INTEGER NOT NULL,
    height            INTEGER NOT NULL,
    params            BLOB,
    prior_focal_length INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS images (
    image_id   INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name       TEXT NOT NULL UNIQUE,
    camera_id  INTEGER NOT NULL,
    prior_qw   REAL,
    prior_qx   REAL,
    prior_qy   REAL,
    prior_qz   REAL,
    prior_tx   REAL,
    prior_ty   REAL,
    prior_tz   REAL
);

CREATE TABLE IF NOT EXISTS keypoints (
    image_id INTEGER PRIMARY KEY NOT NULL,
    rows     INTEGER NOT NULL,
    cols     INTEGER NOT NULL,
    data     BLOB
);

CREATE TABLE IF NOT EXISTS descriptors (
    image_id INTEGER PRIMARY KEY NOT NULL,
    rows     INTEGER NOT NULL,
    cols     INTEGER NOT NULL,
    data     BLOB
);

CREATE TABLE IF NOT EXISTS matches (
    pair_id INTEGER PRIMARY KEY NOT NULL,
    rows    INTEGER NOT NULL,
    cols    INTEGER NOT NULL,
    data    BLOB
);

CREATE TABLE IF NOT EXISTS two_view_geometries (
    pair_id INTEGER PRIMARY KEY NOT NULL,
    rows    INTEGER NOT NULL,
    cols    INTEGER NOT NULL,
    data    BLOB,
    config  INTEGER NOT NULL,
    F       BLOB,
    E       BLOB,
    H       BLOB,
    qvec    BLOB,
    tvec    BLOB
);
"#;

/// Computes COLMAP's canonical pair id for an unordered pair of image ids.
///
/// Must match COLMAP exactly: id1/id2 are ordered smaller-first before encoding.
pub fn image_pair_to_pair_id(id1: i64, id2: i64) -> i64 {
  if id1 > id2 {
    MAX_NUM_IMAGES * id2 + id1
  } else {
    MAX_NUM_IMAGES * id1 + id2
  }
}

pub fn open(path: &Path) -> Result<Connection> {
  let conn = Connection::open(path).with_context(|| format!("opening database at {}", path.display()))?;
  conn.execute_batch(SCHEMA).context("creating database schema")?;
  Ok(conn)
}

/// Looks up the camera_id for a given (width, height), inserting a new SIMPLE_RADIAL
/// camera record if one doesn't already exist for that resolution.
pub fn get_or_create_camera(conn: &Connection, width: u32, height: u32) -> Result<i64> {
  if let Ok(camera_id) = conn.query_row(
    "SELECT camera_id FROM cameras WHERE width = ?1 AND height = ?2 LIMIT 1",
    (width, height),
    |row| row.get::<_, i64>(0),
  ) {
    return Ok(camera_id);
  }

  let f = (width.max(height) as f64) * 1.2;
  let cx = width as f64 / 2.0;
  let cy = height as f64 / 2.0;
  let k = 0.0_f64;
  let params: [f64; 4] = [f, cx, cy, k];
  let params_blob: &[u8] = bytemuck::cast_slice(&params);

  conn
    .execute(
      "INSERT INTO cameras (model, width, height, params, prior_focal_length) VALUES (?1, ?2, ?3, ?4, 0)",
      (SIMPLE_RADIAL_MODEL_ID, width, height, params_blob),
    )
    .context("inserting camera")?;
  Ok(conn.last_insert_rowid())
}

/// Inserts an image row, returning its assigned image_id.
pub fn insert_image(conn: &Connection, name: &str, camera_id: i64) -> Result<i64> {
  conn
    .execute(
      "INSERT INTO images (name, camera_id) VALUES (?1, ?2)",
      (name, camera_id),
    )
    .with_context(|| format!("inserting image row for {name}"))?;
  Ok(conn.last_insert_rowid())
}

/// Returns (width, height) of the camera associated with an image.
pub fn image_dimensions(conn: &Connection, image_id: i64) -> Result<(u32, u32)> {
  conn
    .query_row(
      "SELECT c.width, c.height FROM images i JOIN cameras c ON c.camera_id = i.camera_id WHERE i.image_id = ?1",
      [image_id],
      |row| Ok((row.get::<_, i64>(0)? as u32, row.get::<_, i64>(1)? as u32)),
    )
    .with_context(|| format!("reading dimensions for image_id {image_id}"))
}

pub fn all_images(conn: &Connection) -> Result<Vec<(i64, String)>> {
  let mut stmt = conn.prepare("SELECT image_id, name FROM images ORDER BY name")?;
  let rows = stmt
    .query_map((), |row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)))?
    .collect::<rusqlite::Result<Vec<_>>>()?;
  Ok(rows)
}

/// keypoints.data is float32[N][4]: [x, y, scale, orientation] in original pixel coordinates.
pub fn write_keypoints(conn: &Connection, image_id: i64, keypoints: &[[f32; 4]]) -> Result<()> {
  let data: &[u8] = bytemuck::cast_slice(keypoints);
  conn
    .execute(
      "INSERT INTO keypoints (image_id, rows, cols, data) VALUES (?1, ?2, 4, ?3)",
      (image_id, keypoints.len() as i64, data),
    )
    .context("inserting keypoints")?;
  Ok(())
}

pub fn read_keypoints(conn: &Connection, image_id: i64) -> Result<Vec<[f32; 4]>> {
  let data: Vec<u8> = conn.query_row("SELECT data FROM keypoints WHERE image_id = ?1", [image_id], |row| {
    row.get(0)
  })?;
  let flat: &[f32] = bytemuck::cast_slice(&data);
  Ok(flat.chunks_exact(4).map(|c| [c[0], c[1], c[2], c[3]]).collect())
}

/// descriptors.data is float32[N][128], row-major.
pub fn write_descriptors(
  conn: &Connection,
  image_id: i64,
  rows: usize,
  cols: usize,
  descriptors: &[f32],
) -> Result<()> {
  debug_assert_eq!(descriptors.len(), rows * cols);
  let data: &[u8] = bytemuck::cast_slice(descriptors);
  conn
    .execute(
      "INSERT INTO descriptors (image_id, rows, cols, data) VALUES (?1, ?2, ?3, ?4)",
      (image_id, rows as i64, cols as i64, data),
    )
    .context("inserting descriptors")?;
  Ok(())
}

/// Returns (rows, cols, flat row-major descriptor data).
pub fn read_descriptors(conn: &Connection, image_id: i64) -> Result<(usize, usize, Vec<f32>)> {
  let (rows, cols, data): (i64, i64, Vec<u8>) = conn.query_row(
    "SELECT rows, cols, data FROM descriptors WHERE image_id = ?1",
    [image_id],
    |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
  )?;
  let flat: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
  Ok((rows as usize, cols as usize, flat))
}

/// matches.data is int32[N][2]: [kpt_idx_in_image_A, kpt_idx_in_image_B].
pub fn write_matches(conn: &Connection, pair_id: i64, matches: &[[i32; 2]]) -> Result<()> {
  let data: &[u8] = bytemuck::cast_slice(matches);
  conn
    .execute(
      "INSERT INTO matches (pair_id, rows, cols, data) VALUES (?1, ?2, 2, ?3)",
      (pair_id, matches.len() as i64, data),
    )
    .context("inserting matches")?;
  Ok(())
}

/// Writes a two_view_geometries row. `f_matrix` is the row-major 3x3 fundamental matrix.
/// `config` should be 3 (UNCALIBRATED) unless calibrated intrinsics are available.
pub fn write_two_view_geometry(
  conn: &Connection,
  pair_id: i64,
  config: i64,
  inlier_matches: &[[i32; 2]],
  f_matrix: &[f64; 9],
) -> Result<()> {
  let data: &[u8] = bytemuck::cast_slice(inlier_matches);
  let f_blob: &[u8] = bytemuck::cast_slice(f_matrix.as_slice());
  conn
    .execute(
      "INSERT INTO two_view_geometries (pair_id, rows, cols, data, config, F, E, H, qvec, tvec)
         VALUES (?1, ?2, 2, ?3, ?4, ?5, NULL, NULL, NULL, NULL)",
      (pair_id, inlier_matches.len() as i64, data, config, f_blob),
    )
    .context("inserting two_view_geometries")?;
  Ok(())
}