// Static file server that can plug into the PTTH reverse server

use std::{
	borrow::Cow,
	cmp::min,
	collections::*,
	convert::{Infallible, TryInto},
	error::Error,
	fmt::Debug,
	io::SeekFrom,
	path::{Path, PathBuf},
};

use handlebars::Handlebars;
use percent_encoding::*;
use serde::Serialize;
use tokio::{
	fs::{
		DirEntry,
		File,
		read_dir,
		ReadDir,
	},
	io::AsyncReadExt,
	sync::mpsc::{
		channel,
	},
};
use tracing::instrument;

use regex::Regex;

#[cfg (test)]
use always_equal::test::AlwaysEqual;

#[cfg (not (test))]
use always_equal::prod::AlwaysEqual;

use crate::{
	http_serde::{
		Method,
		Response,
		StatusCode,
	},
	prelude::*,
	prefix_match,
};

#[derive (Serialize)]
struct ServerInfo <'a> {
	server_name: &'a str,
}

#[derive (Serialize)]
struct TemplateDirEntry {
	icon: &'static str,
	trailing_slash: &'static str,
	
	// Unfortunately file_name will allocate as long as some platforms
	// (Windows!) aren't UTF-8. Cause I don't want to write separate code
	// for such a small problem.
	
	file_name: String,
	
	// This could be a Cow with file_name if no encoding was done but
	// it's simpler to allocate.
	
	encoded_file_name: String,
	
	size: Cow <'static, str>,
	
	error: bool,
}

#[derive (Serialize)]
struct TemplateDirPage <'a> {
	#[serde (flatten)]
	server_info: ServerInfo <'a>,
	
	path: Cow <'a, str>,
	entries: Vec <TemplateDirEntry>,
}

fn parse_range_header (range_str: &str) -> (Option <u64>, Option <u64>) {
	use lazy_static::*;
	
	lazy_static! {
		static ref RE: Regex = Regex::new (r"^bytes=(\d*)-(\d*)$").expect ("Couldn't compile regex for Range header");
	}
	
	debug! ("{}", range_str);
	
	let caps = match RE.captures (range_str) {
		Some (x) => x,
		_ => return (None, None),
	};
	let start = caps.get (1).map (|x| x.as_str ());
	let end = caps.get (2).map (|x| x.as_str ());
	
	let start = start.map (|x| u64::from_str_radix (x, 10).ok ()).flatten ();
	
	// HTTP specifies ranges as [start inclusive, end inclusive]
	// But that's dumb and [start inclusive, end exclusive) is better
	
	let end = end.map (|x| u64::from_str_radix (x, 10).ok ().map (|x| x + 1)).flatten ();
	
	(start, end)
}

use std::ops::Range;

#[derive (Debug, PartialEq)]
enum ParsedRange {
	Ok (Range <u64>),
	PartialContent (Range <u64>),
	RangeNotSatisfiable (u64),
}

fn check_range (range_str: Option <&str>, file_len: u64) 
	-> ParsedRange
{
	use ParsedRange::*;
	
	let not_satisfiable = RangeNotSatisfiable (file_len);
	
	let range_str = match range_str {
		None => return Ok (0..file_len),
		Some (x) => x,
	};
	
	let (start, end) = parse_range_header (range_str);
	
	let start = start.unwrap_or (0);
	if start >= file_len {
		return not_satisfiable;
	}
	
	let end = end.unwrap_or (file_len);
	if end > file_len {
		return not_satisfiable;
	}
	if end < start {
		return not_satisfiable;
	}
	
	PartialContent (start..end)
}

async fn read_dir_entry (entry: DirEntry) -> TemplateDirEntry
{
	let file_name = match entry.file_name ().into_string () {
		Ok (x) => x,
		Err (_) => return TemplateDirEntry {
			icon: "⚠️",
			trailing_slash: "",
			file_name: "File / directory name is not UTF-8".into (),
			encoded_file_name: "".into (),
			size: "".into (),
			error: true,
		},
	};
	
	let metadata = match entry.metadata ().await {
		Ok (x) => x,
		Err (_) => return TemplateDirEntry {
			icon: "⚠️",
			trailing_slash: "",
			file_name: "Could not fetch metadata".into (),
			encoded_file_name: "".into (),
			size: "".into (),
			error: true,
		},
	};
	
	let (trailing_slash, icon, size) = {
		let t = metadata.file_type ();
		
		if t.is_dir () {
			("/", "📁", "".into ())
		}
		else {
			let icon = if file_name.ends_with (".mp4") {
				"🎞️"
			}
			else if file_name.ends_with (".avi") {
				"🎞️"
			}
			else if file_name.ends_with (".mkv") {
				"🎞️"
			}
			else if file_name.ends_with (".jpg") {
				"📷"
			}
			else if file_name.ends_with (".jpeg") {
				"📷"
			}
			else if file_name.ends_with (".png") {
				"📷"
			}
			else if file_name.ends_with (".bmp") {
				"📷"
			}
			else {
				"📄"
			};
			
			("", icon, pretty_print_bytes (metadata.len ()).into ())
		}
	};
	
	use percent_encoding::*;
	
	let encoded_file_name = utf8_percent_encode (&file_name, CONTROLS).to_string ();
	
	TemplateDirEntry {
		icon,
		trailing_slash: &trailing_slash,
		file_name,
		encoded_file_name,
		size,
		error: false,
	}
}

async fn serve_root (
	handlebars: &Handlebars <'static>,
) -> Response
{
	let server_info = ServerInfo {
		server_name: "PTTH file server",
	};
	
	let s = handlebars.render ("file_server_root", &server_info).unwrap ();
	let body = s.into_bytes ();
	
	let mut resp = Response::default ();
	resp
		.header ("content-type".to_string (), "text/html".to_string ().into_bytes ())
		.body_bytes (body)
	;
	resp
}

#[instrument (level = "debug", skip (handlebars, dir))]
async fn serve_dir (
	handlebars: &Handlebars <'static>,
	path: Cow <'_, str>,
	mut dir: ReadDir
) -> Response 
{
	let server_info = ServerInfo {
		server_name: "PTTH file server",
	};
	
	let mut entries = vec! [];
	
	while let Ok (Some (entry)) = dir.next_entry ().await {
		entries.push (read_dir_entry (entry).await);
	}
	
	entries.sort_unstable_by (|a, b| a.file_name.partial_cmp (&b.file_name).unwrap ());
	
	let s = handlebars.render ("file_server_dir", &TemplateDirPage {
		path,
		entries,
		server_info,
	}).unwrap ();
	let body = s.into_bytes ();
	
	let mut resp = Response::default ();
	resp
		.header ("content-type".to_string (), "text/html".to_string ().into_bytes ())
		.body_bytes (body)
	;
	resp
}

#[instrument (level = "debug", skip (f))]
async fn serve_file (
	mut f: File,
	should_send_body: bool,
	range: Range <u64>,
	range_requested: bool
) 
	-> Response 
{
	let (tx, rx) = channel (1);
	let body = if should_send_body {
		Some (rx)
	}
	else {
		None
	};
	
	f.seek (SeekFrom::Start (range.start)).await.unwrap ();
	
	info! ("Serving range {}-{}", range.start, range.end);
	
	let content_length = range.end - range.start;
	
	if should_send_body {
		tokio::spawn (async move {
			//println! ("Opening file {:?}", path);
			
			let mut tx = tx;
			let mut bytes_sent = 0;
			let mut bytes_left = content_length;
			
			loop {
				let mut buffer = vec! [0u8; 65_536];
				let bytes_read: u64 = f.read (&mut buffer).await.unwrap ().try_into ().unwrap ();
				
				let bytes_read = min (bytes_left, bytes_read);
				
				buffer.truncate (bytes_read.try_into ().unwrap ());
				
				if bytes_read == 0 {
					break;
				}
				
				if tx.send (Ok::<_, Infallible> (buffer)).await.is_err () {
					warn! ("Cancelling file stream (Sent {} out of {} bytes)", bytes_sent, content_length);
					break;
				}
				
				bytes_left -= bytes_read;
				if bytes_left == 0 {
					debug! ("Finished");
					break;
				}
				
				bytes_sent += bytes_read;
				trace! ("Sent {} bytes", bytes_sent);
				
				//delay_for (Duration::from_millis (50)).await;
			}
		});
	}
	
	let mut response = Response::default ();
	
	response.header (String::from ("accept-ranges"), b"bytes".to_vec ());
	
	if should_send_body {
		if range_requested {
			response.status_code (StatusCode::PartialContent);
			response.header (String::from ("content-range"), format! ("bytes {}-{}/{}", range.start, range.end - 1, range.end).into_bytes ());
		}
		else {
			response.status_code (StatusCode::Ok);
			response.header (String::from ("content-length"), range.end.to_string ().into_bytes ());
		}
		
		response.content_length = Some (content_length);
	}
	else {
		response.status_code (StatusCode::NoContent);
	}
	
	if let Some (body) = body {
		response.body (body);
	}
	
	response
}

fn serve_error (
	status_code: StatusCode,
	msg: &str
)
-> Response
{
	let mut resp = Response::default ();
	resp.status_code (status_code);
	resp.body_bytes (msg.as_bytes ().to_vec ());
	resp
}

fn serve_307 (location: String) -> Response {
	let mut resp = Response::default ();
	resp.status_code (StatusCode::TemporaryRedirect);
	resp.header ("location".to_string (), location.into_bytes ());
	resp.body_bytes (b"Redirecting...".to_vec ());
	resp
}

// Sort of an internal API endpoint to make testing work better.
// Eventually we could expose this as JSON or Msgpack or whatever. For now
// it's just a Rust struct that we can test on without caring about 
// human-readable HTML

#[derive (Debug, PartialEq)]
struct ServeDirParams {
	path: PathBuf,
	dir: AlwaysEqual <ReadDir>,
}

#[derive (Debug, PartialEq)]
struct ServeFileParams {
	send_body: bool,
	range: Range <u64>,
	range_requested: bool,
	file: AlwaysEqual <File>,
}

#[derive (Debug, PartialEq)]
enum InternalResponse {
	Favicon,
	Forbidden,
	MethodNotAllowed,
	NotFound,
	RangeNotSatisfiable (u64),
	Redirect (String),
	Root,
	ServeDir (ServeDirParams),
	ServeFile (ServeFileParams),
}

async fn internal_serve_all (
	root: &Path,
	method: Method,
	uri: &str,
	headers: &HashMap <String, Vec <u8>>,
	hidden_path: Option <&Path>
)
	-> InternalResponse
{
	use InternalResponse::*;
	
	info! ("Client requested {}", uri);
	
	let send_body = match &method {
		Method::Get => true,
		Method::Head => false,
		m => {
			debug! ("Unsupported method {:?}", m);
			return MethodNotAllowed;
		}
	};
	
	if uri == "/favicon.ico" {
		return Favicon;
	}
	
	let uri = match prefix_match ("/files", uri) {
		Some (x) => x,
		None => return Root,
	};
	
	if uri == "" {
		return Redirect ("files/".to_string ());
	}
	
	// TODO: There is totally a dir traversal attack in here somewhere
	
	let encoded_path = &uri [1..];
	
	let path_s = percent_decode (encoded_path.as_bytes ()).decode_utf8 ().unwrap ();
	let path = Path::new (&*path_s);
	
	let mut full_path = PathBuf::from (root);
	full_path.push (path);
	
	debug! ("full_path = {:?}", full_path);
	
	if let Some (hidden_path) = hidden_path {
		if full_path == hidden_path {
			return Forbidden;
		}
	}
	
	let has_trailing_slash = path_s.is_empty () || path_s.ends_with ("/");
	
	if let Ok (dir) = read_dir (&full_path).await {
		if ! has_trailing_slash {
			return Redirect (format! ("{}/", path.file_name ().unwrap ().to_str ().unwrap ()));
		}
		
		let dir = dir.into ();
		
		ServeDir (ServeDirParams {
			dir,
			path: full_path,
		})
	}
	else if let Ok (file) = File::open (&full_path).await {
		let file_md = file.metadata ().await.unwrap ();
		let file_len = file_md.len ();
		
		let range_header = headers.get ("range").map (|v| std::str::from_utf8 (v).ok ()).flatten ();
		
		let file = file.into ();
		
		match check_range (range_header, file_len) {
			ParsedRange::RangeNotSatisfiable (file_len) => RangeNotSatisfiable (file_len),
			ParsedRange::Ok (range) => ServeFile (ServeFileParams {
				file, 
				send_body, 
				range,
				range_requested: false,
			}),
			ParsedRange::PartialContent (range) => ServeFile (ServeFileParams {
				file,
				send_body, 
				range,
				range_requested: true,
			}),
		}
	}
	else {
		NotFound
	}
}

#[instrument (level = "debug", skip (handlebars, headers))]
pub async fn serve_all (
	handlebars: &Handlebars <'static>,
	root: &Path,
	method: Method,
	uri: &str,
	headers: &HashMap <String, Vec <u8>>,
	hidden_path: Option <&Path>
)
	-> Response
{
	use InternalResponse::*;
	
	match internal_serve_all (root, method, uri, headers, hidden_path).await {
		Favicon => serve_error (StatusCode::NotFound, ""),
		Forbidden => serve_error (StatusCode::Forbidden, "403 Forbidden"),
		MethodNotAllowed => serve_error (StatusCode::MethodNotAllowed, "Unsupported method"),
		NotFound => serve_error (StatusCode::NotFound, "404 Not Found"),
		RangeNotSatisfiable (file_len) => {
			let mut resp = Response::default ();
			resp.status_code (StatusCode::RangeNotSatisfiable)
			.header ("content-range".to_string (), format! ("bytes */{}", file_len).into_bytes ());
			resp
		},
		Redirect (location) => serve_307 (location),
		
		Root => serve_root (handlebars).await,
		ServeDir (ServeDirParams {
			path,
			dir,
		}) => serve_dir (handlebars, path.to_string_lossy (), dir.into_inner ()).await,
		ServeFile (ServeFileParams {
			file,
			send_body, 
			range, 
			range_requested,
		}) => serve_file (file.into_inner (), send_body, range, range_requested).await,
	}
}

pub fn load_templates () 
-> Result <Handlebars <'static>, Box <dyn Error>>
{
	let mut handlebars = Handlebars::new ();
	handlebars.set_strict_mode (true);
	
	for (k, v) in vec! [
		("file_server_dir", "file_server_dir.html"),
		("file_server_root", "file_server_root.html"),
	].into_iter () {
		handlebars.register_template_file (k, format! ("ptth_handlebars/{}", v))?;
	}
	
	Ok (handlebars)
}

fn pretty_print_bytes (b: u64) -> String {
	if b < 1024 {
		format! ("{} B", b)
	}
	else if (b + 512) < 1024 * 1024 {
		format! ("{} KiB", (b + 512) / 1024)
	}
	else if (b + 512 * 1024) < 1024 * 1024 * 1024 {
		format! ("{} MiB", (b + 512 * 1024) / 1024 / 1024)
	}
	else {
		format! ("{} GiB", (b + 512 * 1024 * 1024) / 1024 / 1024 / 1024)
	}
}

#[cfg (test)]
mod tests {
	use std::{
		ffi::OsStr,
		path::{
			Component, 
			Path, 
			PathBuf
		},
	};
	
	use maplit::*;
	use tokio::runtime::Runtime;
	
	use always_equal::test::AlwaysEqual;
	
	use crate::http_serde::{
		StatusCode,
	};
	
	#[test]
	fn parse_range_header () {
		for (input, expected) in vec! [
			("", (None, None)),
			("bytes=0-", (Some (0), None)),
			("bytes=0-999", (Some (0), Some (1000))),
			("bytes=111-999", (Some (111), Some (1000))),
		].into_iter () {
			let actual = super::parse_range_header (input);
			assert_eq! (actual, expected);
		}
		
		use super::ParsedRange::*;
		
		for (header, file_len, expected) in vec! [
			(None, 0, Ok (0..0)),
			(None, 1024, Ok (0..1024)),
			
			(Some (""), 0, RangeNotSatisfiable (0)),
			(Some (""), 1024, PartialContent (0..1024)),
			
			(Some ("bytes=0-"), 1024, PartialContent (0..1024)),
			(Some ("bytes=0-999"), 1024, PartialContent (0..1000)),
			(Some ("bytes=0-1023"), 1024, PartialContent (0..1024)),
			(Some ("bytes=111-999"), 1024, PartialContent (111..1000)),
			(Some ("bytes=111-1023"), 1024, PartialContent (111..1024)),
			(Some ("bytes=200-100"), 1024, RangeNotSatisfiable (1024)),
			
			(Some ("bytes=0-"), 512, PartialContent (0..512)),
			(Some ("bytes=0-1023"), 512, RangeNotSatisfiable (512)),
			(Some ("bytes=1000-1023"), 512, RangeNotSatisfiable (512)),
		].into_iter () {
			let actual = super::check_range (header, file_len);
			assert_eq! (actual, expected);
		}
	}
	
	#[test]
	fn pretty_print_bytes () {
		for (input_after, expected_before, expected_after) in vec! [
			(1, "0 B", "1 B"),
			(1024, "1023 B", "1 KiB"),
			(1024 + 512, "1 KiB", "2 KiB"),
			(1023 * 1024 + 512, "1023 KiB", "1 MiB"),
			((1024 + 512) * 1024, "1 MiB", "2 MiB"),
			(1023 * 1024 * 1024 + 512 * 1024, "1023 MiB", "1 GiB"),
			((1024 + 512) * 1024 * 1024, "1 GiB", "2 GiB"),
			
		].into_iter () {
			let actual = super::pretty_print_bytes (input_after - 1);
			assert_eq! (&actual, expected_before);
			
			let actual = super::pretty_print_bytes (input_after);
			assert_eq! (&actual, expected_after);
		}
	}
	
	#[test]
	fn i_hate_paths () {
		let mut components = Path::new ("/home/user").components ();
		
		assert_eq! (components.next (), Some (Component::RootDir));
		assert_eq! (components.next (), Some (Component::Normal (OsStr::new ("home"))));
		assert_eq! (components.next (), Some (Component::Normal (OsStr::new ("user"))));
		assert_eq! (components.next (), None);
		
		let mut components = Path::new ("./home/user").components ();
		
		assert_eq! (components.next (), Some (Component::CurDir));
		assert_eq! (components.next (), Some (Component::Normal (OsStr::new ("home"))));
		assert_eq! (components.next (), Some (Component::Normal (OsStr::new ("user"))));
		assert_eq! (components.next (), None);
		
		let mut components = Path::new (".").components ();
		
		assert_eq! (components.next (), Some (Component::CurDir));
		assert_eq! (components.next (), None);
	}
	
	#[test]
	fn file_server () {
		use crate::{
			http_serde::Method,
			//prelude::*,
		};
		use super::{
			InternalResponse,
			internal_serve_all,
			load_templates,
			serve_all,
			ServeDirParams,
			ServeFileParams,
		};
		
		tracing_subscriber::fmt ().try_init ().ok ();
		let mut rt = Runtime::new ().unwrap ();
		
		rt.block_on (async {
			let handlebars = load_templates ().unwrap ();
			let file_server_root = PathBuf::from ("./");
			let headers = Default::default ();
			
			for (uri_path, expected_status) in vec! [
				("/", StatusCode::Ok),
				("/files", StatusCode::TemporaryRedirect),
				("/files/src", StatusCode::TemporaryRedirect),
				("/files/src/", StatusCode::Ok),
			].into_iter () {
				let resp = serve_all (
					&handlebars, 
					&file_server_root,
					Method::Get,
					uri_path, 
					&headers,
					None
				).await;
				
				assert_eq! (resp.parts.status_code, expected_status);
			}
			
			{
				use InternalResponse::*;
				
				for (uri_path, expected) in vec! [
					("/", Root),
					("/files", Redirect ("files/".to_string ())),
					("/files/src", Redirect ("src/".to_string ())),
					("/files/src/bad_passwords.txt", ServeFile (ServeFileParams {
						send_body: true,
						range: 0..1_048_576,
						range_requested: false,
						file: AlwaysEqual::testing_blank (),
					})),
				].into_iter () {
					let resp = internal_serve_all (
						&file_server_root,
						Method::Get,
						uri_path, 
						&headers,
						None
					).await;
					
					assert_eq! (resp, expected);
				}
				
				let resp = internal_serve_all (
					&file_server_root,
					Method::Get,
					"/files/src/bad_passwords.txt", 
					&hashmap! {
						"range".into () => b"bytes=0-2000000".to_vec (),
					},
					None
				).await;
				
				assert_eq! (resp, RangeNotSatisfiable (1_048_576));
			}
		});
	}
}