//! Persistent store of all calculated image descriptions.
//! Can be queried to find identical or similar images.

use std::collections::{HashMap, HashSet};
use crate::descriptors::Descriptor;
use std::fs;
use std::path::Path;
use std::fmt;
use log::{info, debug, error};
use bktree::*;

#[derive(Debug)]
pub enum SaveError {
	Serialization,
	File,
}

/// Uses a hashmap to map descriptors to buckets of files.
/// Also keeps a BK-tree for quick distance ranking
pub struct DescriptorStore<T: Descriptor> {
	descriptor: T,
	/// Main hashmap that maps descriptors to buckets of filenames
	map: HashMap<u64, Vec<String>>, 

	/// Place to load and store the map on file
	save_location: std::path::PathBuf, 

	/// BK-tree for fast nearest neighbour
	bktree: BkTree<u64>, 

	/// Cache for seen files to skip on initial load
	seen: Option<HashSet<String>>, 
}

impl<T: Descriptor> DescriptorStore<T> {
	/// Makes a new empty DescriptorStore with default settings
	pub fn new(descriptor: T) -> Self {
		let map: HashMap<u64, Vec<String>> = HashMap::new();
		//let seen: HashSet<String> = HashSet::new();
		//let seen = None;
		DescriptorStore {
			map,
			descriptor,
			save_location: std::path::PathBuf::from("store.messagepack"),
			bktree: BkTree::new(hamming_distance),
			seen: None,
		}
	}

	/// Sets the file location and loads data from file (if available)
	pub fn with_file<P: AsRef<Path>>(mut self, path: P) -> Self {
		self.save_location = std::path::PathBuf::from(path.as_ref());
		let map_file = fs::read(&self.save_location);
		match map_file {
			Ok(f) => {
				self.map = rmp_serde::from_slice(&f).unwrap();
				for i in self.map.keys() {
					self.bktree.insert(*i);
				}
				let mut seen: HashSet<String> = HashSet::new();
				for bucket in self.map.values() {
					for element in bucket {
						seen.insert(element.clone());
					}
				}
				self.seen = Some(seen);
			},
			Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
		};
		self
	}

	pub fn save(&self) -> Result<(), SaveError> {
		let serialized: Vec<u8> = match rmp_serde::to_vec(&self.map) {
			Ok(value) => value,
			Err(_e) => return Err(SaveError::Serialization),
		};
		match fs::write(&self.save_location, &serialized) {
			Ok(_) => Ok(()),
			Err(_e) => Err(SaveError::File),
		}
	}

	/// Returns true iff the store already contains the key
	pub fn contains(&self, key: u64) -> bool {
		self.map.contains_key(&key)
	}

	/// Returns true iff the store already contains the value in some bucket. 
	/// Will use a hashmap cache if available
	pub fn has_value(&self, value: &String) -> bool {
		match &self.seen {
			Some(set) => {
				set.contains(value)
			},
			None => {
				for bucket in self.map.values() {
					if bucket.contains(&value) {
						return true;
					}
				}
				false
			}
		}
	}

	/// Inserts a single value into the store
	pub fn insert(&mut self, key: &u64, value: String) {
		let bucket = match self.map.get(&key) {
			Some(b) => {
				let mut n = b.clone();
				if !n.contains(&value) {
					n.push(value);
				}
				n
			},
			None => vec![value],
		};
		self.map.insert(*key, bucket);
		self.bktree.insert(*key);
	}

	/// Calculates all descriptions with a given descriptor 
	pub fn insert_directory<U: AsRef<Path>>(&mut self, dir: U) {
		for node in fs::read_dir(dir).unwrap() {
			let file = node.expect("Error walking directory");
			let name = match file.file_name().into_string() {
				Ok(v) => v,
				Err(e) => {
					error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e);
					continue
				}
			};
			if !self.has_value(&name) { // !self.contains(name.to_string()) {
				info!("Processing {}", name);
				let img = match image::open(file.path()) {
					Ok(v) => v,
					Err(e) => {
						error!("Failed to process {}: {}", name, e);
						continue
					}
				};
				let phash = self.descriptor.describe(&img);
				if self.contains(phash) {
					println!("{} duplicate of {:?}", name, self.map.get(&phash));
				}
				self.insert(&phash, name);
			} else {
				debug!("{} already known, skipping.", name);
			}
			self.save().expect("error");
		}
		// We have done our bulk loading, so we can unset our hashset cache:
		self.seen = None;
	}

	/// Nearest neighbours
	pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
		let mut neighbours = self.bktree.find(from, max_distance);
		neighbours.sort_by(|a, b| a.1.cmp(&b.1));
		neighbours
	}

	pub fn print_nn(&self, from: u64, max_distance: isize, show_images: bool) {
		//println!("{from:064b}");
		let neighbours = self.nn(from, max_distance);
		let (term_width, _) = viuer::terminal_size();
		let mut x = 0;
		let mut y = 8;
		for (element, distance) in neighbours {
			let elements = self.map.get(element);
			match elements {
				Some(paths) => {
					for path in paths {
						println!("{element:064b}: {:?} (distance: {distance})", path);
						if show_images {
							if x+16 >= term_width {
								x = 0;
								y += 8;
							}
							let conf = viuer::Config {
								width: Some(16),
								height: Some(8),
								x,
								y,
								use_kitty: false,
								..Default::default()
							};
							x += 16;
							let path = "data/".to_string() + path;
							let img = image::open(&path).unwrap();
        					let img = img.grayscale().thumbnail_exact(8, 8);
							viuer::print(&img, &conf).expect("Image printing failed.");
						}
					}
				},
				None => ()
			};
		}
	}

	pub fn print_most_dups(&self) {
		let mut most = 0;
		for bucket in self.map.values() {
			let len = bucket.len();
			if len > most {
				most = len;
			}
		}

		let mut output = String::new();
		for (key, bucket) in self.map.iter() {
			if bucket.len() == most {
				let key_entry = format!("{key:064b}:\n\t\t");
				output.push_str(&key_entry);
				for value in bucket {
					let value_entry = format!("{value}\n");
					output.push_str(&value_entry);
				}
				output.push_str("\n");
			}
		}
		print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
	}
}

impl<T: Descriptor> fmt::Display for DescriptorStore<T> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		let mut output = String::new();
		for (key, bucket) in self.map.iter() {
			let key_entry = format!("{key:064b}:\n");
			output.push_str(&key_entry);
			for value in bucket {
				let value_entry = format!("\t{value}\n");
				output.push_str(&value_entry);
			}
			output.push_str("\n");
		}
		write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
	}
}