mod parse;
mod crypto;
mod upload;
mod thumbnailing;
mod protobuf;
mod errors;

use std::{borrow::Cow, io::stdout};
use std::io::{Read, Write, self};
use std::fs::{File, self, OpenOptions};
use std::path::Path;

use clap::Parser;
use errors::AviaryUploadError;
use itertools::{Itertools, Either};
use parse::{CreateArgs, Command, DownloadArgs};
use ::protobuf::Message;
use sanitise_file_name::sanitise_with_options;

use crate::errors::AviaryDownloadError;
use crate::protobuf::image::Format;

fn trim_url<'a>(base_url: &str, url: &'a str) -> Option<&'a str> {
	if url.starts_with(base_url) {
		let shortened = url
			.trim()
			.trim_start_matches(base_url)
			.trim_matches('/')
			.trim_end_matches(".bin");
		if shortened.len() > 50 {
			None
		} else {
			Some(shortened)
		}
	} else {
		None
	}
}

fn main() {
	let args = parse::Args::parse();
	let server_no_trailing_slash = args.server.trim_end_matches('/');
	let full_server = if server_no_trailing_slash.starts_with("http") {
		Cow::Borrowed(server_no_trailing_slash)
	} else {
		Cow::Owned(format!("https://{}", server_no_trailing_slash))
	};
	match args.command {
		Command::Create(create_args) => create(&*full_server, create_args),
		Command::Download(download_args) =>
			if let Err(e) = download(&*full_server, download_args) {
				println!("{e}")
			},
	}
}

fn create(server: &str, args: CreateArgs) {
	print!("Checking files...");
	let (files, errors): (Vec<_>, Vec<_>) = args.images.iter()
		.map(Path::new)
		.map(|path| File::open(path)
			 .map(|file| (path, file))
			 .map_err(|e| AviaryUploadError::from_open_error(e.kind(), &path))
		)
		.partition_result();

	if !errors.is_empty() {
		println!(" \x1b[31mError!\x1b[0m");
		let (nonexistant, noread): (Vec<_>, Vec<_>) = errors.iter()
			.partition_map(|e| match e {
				AviaryUploadError::FileDNE(path) => Either::Left(path),
				AviaryUploadError::ReadPermissionDenied(path) => Either::Right(path),
				other => panic!("This error should not be possible! {other:?}")
			});
		if !nonexistant.is_empty() {
			println!("\nWe didn't see any files at the following locations:");
			nonexistant.iter().for_each(|path| println!("\x1b[37m- \x1b[31m{}", path.display()));
		}
		if !noread.is_empty() {
			println!("\x1b[0m\nWe found these files, but didn't have permission to open them:");
			noread.iter().for_each(|path| println!("\x1b[37m- \x1b[31m{}", path.display()));
		}
	} else {
		println!(" \x1b[32mDone!\n");
		let agent = upload::get_agent();
		let index_url = files.into_iter()
			.inspect(|(path, _)| print!("\x1b[36m{}\x1b[0m\n\x1b[37m├─\x1b[0m Reading...      ", path.display()))
			.inspect(|_| drop(stdout().flush()))
			.map(|(path, mut file)|
				 (|| {
					let mut buff = Vec::with_capacity(file.metadata()?.len() as usize);
					file.read_to_end(&mut buff)?;
					Ok((path, buff))
				})().map_err(|e| AviaryUploadError::StreamReadError(path.to_owned(), e))
			)
			.inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Thumbnailing... ") })
			.inspect(|_| drop(stdout().flush()))
			.map(|r| r.and_then(|(path, raw_dat)| {
				let (full, w, h, thumbnail, blurhash, format) = thumbnailing::thumbnail(&raw_dat)
					.map_err(|_| AviaryUploadError::ImageFormatError(path.to_owned()))?;
				Ok((full.map(Either::Right).unwrap_or(Either::Left(raw_dat)), w, h, thumbnail, blurhash, format))
			}))
			.inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Encrypting...   ")})
			.inspect(|_| drop(stdout().flush()))
			.map_ok(|(raw_dat, w, h, thumbnail, blurhash, format)| {
				let key = crypto::make_key();
				(
					key,
					w, h,
					crypto::encrypt(&key, &crypto::NONCE_A, &raw_dat),
					crypto::encrypt(&key, &crypto::NONCE_B, &thumbnail),
					blurhash,
					format
				)
			})
			.inspect(|r| if r.is_ok() { print!("\x1b[32mDone!\n\x1b[37m└─\x1b[0m Uploading...    ")})
			.inspect(|_| drop(stdout().flush()))
			.map(|r| r.and_then(|(key, w, h, full_img, thumb, blurhash, format)|
				upload::put_data(&agent, server, &thumb)
					.and_then(|thumb_url|
						  upload::put_data(&agent, server, &full_img)
							.map(|full_url| (key, w, h, full_url, thumb_url, blurhash, format)))
					.map_err(AviaryUploadError::from_upload_error)
			))
			.map(|r| r.and_then(|(key, w, h, full_url, thumb_url, blurhash, format)| {
				let full_trimmed = trim_url(server, &full_url);
				let thmb_trimmed = trim_url(server, &thumb_url);
				if let (Some(full_url), Some(thmb_url)) = (full_trimmed, thmb_trimmed) {
					Ok((key, w, h, full_url.to_owned(), thmb_url.to_owned(), blurhash, format))
				} else {
					Err(AviaryUploadError::ServerError(format!("Received bad response from server: {}", full_url)))
				}
			}))
			.inspect(|r| if r.is_ok() { println!("\x1b[32mDone!\n")})
			.map_ok(|(key, width, height, full_url, thumb_url, blurhash, format)| protobuf::image::Image {
				key: key.into(),
				format: format.into(),
				width, height, full_url, thumb_url, blurhash,
				special_fields: Default::default()
			})
			.collect::<Result<Vec<_>, _>>()
			.and_then(|image_info|{
				let index = protobuf::index::Index {
					images: image_info,
					title: args.title,
					desc: None,
					special_fields: Default::default(),
				};
				let index_key = crypto::make_key();
				let encrypted_index = crypto::encrypt(
					&index_key,
					&crypto::NONCE_A,
					&index.write_to_bytes()
						.expect("Error serializing protocol buffers")
				);
				let encoded_key = base64::encode(index_key);
				print!("\x1b[0mUploading index... ");
				upload::put_data(&agent, server, &encrypted_index)
					.map_err(|e| AviaryUploadError::from_upload_error(e))
					.map(|url| format!("{}#{}", url.trim().trim_end_matches(".bin"), &encoded_key))
			});
		match index_url {
			Ok(url) =>
				println!("\x1b[32mDone!\n\n\x1b[34mYour gallery is: \x1b[1;0m{}", url),
			Err(e) =>
				print!("{}", e),
		}
	}
}

const SANITIZATION_OPTIONS: sanitise_file_name::Options<fn(char) -> Option<char>> = sanitise_file_name::Options {
	length_limit: 127,
	reserve_extra: 0,
	extension_cleverness: false,
	most_fs_safe: true,
	windows_safe: true,
	url_safe: false,
	normalise_whitespace: false,
	trim_spaces_and_full_stops: true,
	trim_more_punctuation: false,
	remove_control_characters: true,
	remove_reordering_characters: false,
	replace_with: |_| Some('_'),
	collapse_replacements: true,
	six_measures_of_barley: "Unnamed-Album"
};

fn download(server: &str, args: DownloadArgs) -> Result<(), AviaryDownloadError> {
	let mut download_buffer = Vec::with_capacity(5_000_000);
	let mut decrypt_buffer = Vec::with_capacity(5_000_000);
	let index_url = format!("{}/{}.bin", server, args.id);
	let key: [u8; 32] = base64::decode(args.key)
		.map_err(|_| AviaryDownloadError::MalformedKey(true))?
		.try_into()
		.map_err(|_| AviaryDownloadError::MalformedKey(false))?;
	let download_agent = upload::get_agent();

	print!("Downloading index... ");
	let encrypted_index = upload::get_data(&download_agent, &index_url, 500_000, &mut download_buffer)
		.map_err(AviaryDownloadError::from_download_error(true))?;
	let serialized_index = crypto::decrypt(&key, &crypto::NONCE_A, &encrypted_index, &mut decrypt_buffer)
		.ok_or(AviaryDownloadError::KeyMismatch)?;
	let index = protobuf::index::Index::parse_from_bytes(&serialized_index)
		.map_err(|_| AviaryDownloadError::CorruptIndex("Index was not a valid protobuf structure"))?;
	let image_count = index.images.len();
	println!(
		"\x1b[32mDone!\x1b[0m\n\nGallery Name: {}\n Image Count: {}\n\n(desc not implemented)\n",
		index.title.as_ref().map(String::as_str).unwrap_or("<not set>"),
		image_count
	);

	let dest_dir: Cow<Path> = args.output.map(Cow::Owned)
		.unwrap_or_else(||
			index.title.map(|title|
				Cow::Owned(sanitise_with_options(&title, &SANITIZATION_OPTIONS).into()))
			.unwrap_or(Cow::Borrowed("Unnamed-Album".as_ref())));

	fs::create_dir_all(&dest_dir).map_err(|e| match e.kind() {
		io::ErrorKind::PermissionDenied =>
			AviaryDownloadError::PermissionDenied(dest_dir.as_ref().to_owned()),
		_ => AviaryDownloadError::FilesystemError(e, dest_dir.as_ref().to_owned()),
	})?;
	for (indx, image) in index.images.into_iter().enumerate() {

		let extension = image.format.enum_value().map(|f| match f {
			Format::PNG => "png",
			Format::WEBP => "webp",
			Format::AVIF => "avif",
			Format::JPG => "jpg",
			Format::GIF => "gif",
		}).unwrap_or("unk");

		// Print a little blurhash
		const BH_WIDTH: u32 = 12;
		const BH_HEIGHT: u32 = 5;
		let blur = blurhash::decode(&image.blurhash, BH_WIDTH, BH_HEIGHT, 1.0);
		let mut blur = blur.iter();
		for y in 0..BH_HEIGHT {
			for _ in 0..BH_WIDTH {
				let r = blur.next().unwrap();
				let g = blur.next().unwrap();
				let b = blur.next().unwrap();
				let _ = blur.next().unwrap();
				print!("\x1b[38;2;{r};{g};{b}m█");
			}
			match y {
				1 => println!("\x1b[0m Image {}/{}", indx + 1, image_count),
				2 => println!("\x1b[0m Format: {}", extension),
				_ => println!(""),
			}
		}
		print!("\x1b[37m├─\x1b[0m Reserving file...       ");
		stdout().flush().unwrap();
		let path = dest_dir.join(format!("{indx:03}.{extension}"));
		let mut dest_file =  OpenOptions::new()
			.write(true)
			.truncate(true)
			.create(true)
			.create_new(!args.force)
			.open(&path)
			.map_err(|e| match e.kind() {
				io::ErrorKind::PermissionDenied =>
					AviaryDownloadError::PermissionDenied(path.clone()),
				io::ErrorKind::AlreadyExists =>
					AviaryDownloadError::NotOverwriting(path.clone()),
				io::ErrorKind::Interrupted =>
					AviaryDownloadError::Cancelled,
				_ => AviaryDownloadError::FilesystemError(e, path.clone())
			})?;
		let key = image.key.try_into()
			.map_err(|_|
				AviaryDownloadError::CorruptIndex("Key of invalid length specified in index")
			)?;
		print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Dowloading...	   ");
		stdout().flush().unwrap();
		let encrypted_thumbnail = upload::get_data(
			&download_agent,
			&format!("{}/{}.bin", server, image.full_url),
			7680_000,
			&mut download_buffer
		).map_err(AviaryDownloadError::from_download_error(false))?;
		print!("\x1b[32mDone!\n\x1b[37m├─\x1b[0m Decrypting...	   ");
		stdout().flush().unwrap();
		let decrypted_thumb = crypto::decrypt(
			&key,
			&crypto::NONCE_B,
			encrypted_thumbnail,
			&mut decrypt_buffer
		).ok_or(AviaryDownloadError::ExpiredImage)?;
		print!("\x1b[32mDone!\n\x1b[37m└─\x1b[0m Saving...		   ");
		stdout().flush().unwrap();
		dest_file.write_all(&decrypted_thumb)
			.map_err(|e| match e.kind() {
				io::ErrorKind::PermissionDenied =>
					AviaryDownloadError::PermissionDenied(path.clone()),
				io::ErrorKind::Interrupted =>
					AviaryDownloadError::Cancelled,
				_ => AviaryDownloadError::FilesystemError(e, path.clone())
			})?;
		println!("\x1b[32mDone!\n");
	}

	Ok(())
}