mod types;
mod html;
mod api;

use {
	std::{
		sync::Arc,
		net::SocketAddr,
		collections::HashMap,
		process::exit,
		time::{
			SystemTime,
			UNIX_EPOCH,
		},
	},
	chrono::{
		DateTime,
	},
	http_body_util::{
		Full,
		BodyExt,
	},
	hyper::{
		StatusCode,
		Request,
		Response,
		Method,
		HeaderMap,
		header::HeaderValue,
		body::{
			Bytes,
			Incoming,
		},
		server::conn::http1 as Server,
		service::service_fn,
	},
	hyper_util::{
		rt::TokioIo,
	},
	tokio::{
		net::TcpListener,
	},
	rsa::{
		pkcs1::{
			EncodeRsaPublicKey,
			EncodeRsaPrivateKey,
		},
		RsaPrivateKey,
		RsaPublicKey,
	},
	urlencoding::{
		decode as url_decode,
	},
	jsonwebtoken as jwt,
	jwt::{
		Header,
		Algorithm,
		TokenData,
		Validation,
		EncodingKey,
		DecodingKey,
		encode as jwt_encode,
		decode as jwt_decode,
	},
	serde_json::{
		Value as Json,
		json,
	},
	uuid::{
		Uuid,
	},
	skytable::{
		query,
		Config,
		pool::{
			self,
			ConnectionMgrTcp
		},
		error::Error as SkyError,
	},
	bb8::{
		Pool,
	},
	crate::{
		types::{
			users::Users,
			sites::Sites,
		},
		html::*,
	},
};

type Res<T, E> = std::result::Result<T, E>;
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
type DBPool = Arc<Pool<ConnectionMgrTcp>>;
type FullBytes = Result<Response<Full<Bytes>>>;


const PORT: u16 = 8083;

const DB_POOL: u32 = 32;
const DB_ADDR: &str = "192.168.1.49";
const DB_PORT: u16 = 2003;
const DB_USER: &str = "root";
const DB_PASS: &str = "rootpass12345678";

const TOKEN_LIFETIME: u32 = 300;
const REFRESH_LIFETIME: u32 = 2_678_400;


#[tokio::main]
async fn main() -> Result<()> {
	println!("Starting.");
	
	print!("Binding port...");
	let addr = SocketAddr::from(([0, 0, 0, 0], PORT));
	let listener = TcpListener::bind(addr).await;
	if listener.is_err() {
		println!(" Error.");
		exit(1);
	}
	let listener = listener?;
	println!(" OK.");

	print!("Connecting to DB...");
	let pool = pool::get_async(DB_POOL, Config::new(DB_ADDR, DB_PORT, DB_USER, DB_PASS)).await;
	let pool = Arc::new(pool.unwrap());

	if pool.get().await.is_err() {
		println!(" Error.");
		exit(1);
	}
	println!(" OK.");

	init_tables(pool.clone()).await?;

	println!("Server started on port: {}", PORT);

	loop {
		let (stream, _) = listener.accept().await?;

		let ip = format!("{:?}", stream.peer_addr().unwrap().ip());

		let io = TokioIo::new(stream);

		let this_pool = pool.clone();
		tokio::task::spawn(async move {
			if let Err(err) = Server::Builder::new().serve_connection(
				io, service_fn(move |req| handle_connection(
					req,
					this_pool.clone(),
					ip.clone()
				))
			).await {
				println!("Error serving connection: {:?}", err);
			}
		});
	}
}

async fn handle_connection(req: Request<Incoming>, pool: DBPool, ip: String) -> FullBytes {
	let t = time_ns();

	if ip == "1.1.1.1" {
	}

	let res = Response::new(Full::new(Bytes::new()));
	let (mut parts, _) = res.into_parts();
	let body: String;
	let restype: HeaderValue;

	let mut headers = parts.headers.clone();

	let cookies = get_cookies(req.headers().clone());
	let token = cookies.get("token");
	let token = match token.is_none() {
		false => token.unwrap(),
		_ => ""
	};

	match <str as AsRef<str>>::as_ref(req.uri().path()) {
		x if x.starts_with("/api/") => {}
		_ => {
//			println!("{}", token);

			if token != "" && jwt_verify(pool.clone(), token)
				.await?.claims.as_object().unwrap().len() == 0
			{
				println!("Invalid suka");
				parts.status = StatusCode::FOUND;
				set_cookie(&mut headers, "token", "");
				set_location(&mut headers, "/");
				parts.headers = headers;
				return Ok(Response::from_parts(parts, Full::new(Bytes::new())));
			}
		}
	}

	(body, parts.status, restype) = match req.uri().path().as_ref() {
		"/" => uri_index(),
		"/cabinet" => uri_login(req, pool.clone(), &mut headers).await?,
		"/login" => uri_login(req, pool.clone(), &mut headers).await?,
		"/register" => uri_register(req, pool.clone(), &mut headers).await?,
		"/recover" => uri_recover(),
		x if x.starts_with("/@") => uri_user(req, pool.clone()).await?,
		x if x.starts_with("/api/") => api::endpoint(req, pool.clone()).await,
		_ => uri_404()
	};

	headers.insert(hyper::header::CONTENT_TYPE, restype);
	parts.headers = headers;

	let body = body.replace("RENDER_TIME", &format!("{}", time_ns() - t));
	Ok(Response::from_parts(parts, Full::new(Bytes::from(body))))
}

fn build_html(body: &str) -> String {
	format!("{}{}{}{}", HEADER_HTML, CSS3, body, FOOTER_HTML)
}

fn set_cookie(headers: &mut HeaderMap, key: &str, value: &str) {
	let time = DateTime::from_timestamp((time() + REFRESH_LIFETIME) as i64, 0)
		.expect("REASON")
		.to_rfc2822();
	let time = time.replace("+0000", "GMT");
	headers.append(
		hyper::header::SET_COOKIE,
		HeaderValue::from_str(format!("{}={}; HttpOnly; Expires={}", key, value, time).as_str())
			.unwrap()
	);
}

fn get_cookies(headers: HeaderMap) -> HashMap<String, String> {
	let header = headers.get(hyper::header::COOKIE);
	let cookies = match header.is_none() {
		false => header.unwrap().to_str().unwrap(),
		_ => ""
	};

	double_split(cookies.to_owned(), ";", "=")
}

fn set_location(headers: &mut HeaderMap, location: &str) {
	headers.append(
		hyper::header::LOCATION,
		HeaderValue::from_str(format!("{}", location).as_str()).unwrap()
	);
}

async fn uri_login(req: Request<Incoming>, pool: DBPool, headers: &mut HeaderMap) -> Result<(String, StatusCode, HeaderValue)> {
	if *req.method() == Method::POST {
		let body = get_body_from_request(req).await?;
		let body = String::from_utf8(body).unwrap();
		let body = double_split(body, "&", "=");

		let (access, refresh) = login_user(pool.clone(), body).await?;

		set_cookie(headers, "token", &access);
		set_cookie(headers, "refresh", &refresh);
	}

	let restype: HeaderValue = "text/html".parse().unwrap();
	Ok((build_html(LOGIN_HTML), StatusCode::OK, restype))
}

fn uri_index() -> (String, StatusCode, HeaderValue) {
	let restype: HeaderValue = "text/html".parse().unwrap();
	(build_html(INDEX_HTML), StatusCode::OK, restype)
}

async fn uri_register(req: Request<Incoming>, pool: DBPool, headers: &mut HeaderMap) -> Result<(String, StatusCode, HeaderValue)> {
	let mut body = "".to_owned();
	let mut status = StatusCode::OK;
	let restype: HeaderValue = "text/html".parse().unwrap();

	if *req.method() == Method::POST {
		let request = get_body_from_request(req).await?;
		let request = String::from_utf8(request).unwrap();
		let request = double_split(request, "&", "=");

		match create_user(pool.clone(), request).await? {
			true => {
				println!("Created");
				set_location(headers, "/login");
			},
			_ => {
				println!("Failed");
				set_location(headers, "/register");
			}
		}

		status = StatusCode::FOUND;
	}
	else {
		body = build_html(REG_HTML);
	}

	Ok((body, status, restype))
}

fn uri_recover() -> (String, StatusCode, HeaderValue) {
	let restype: HeaderValue = "text/html".parse().unwrap();
	(build_html(RECOVER_HTML), StatusCode::OK, restype)
}

async fn uri_user(req: Request<Incoming>, pool: DBPool) -> Res<(String, StatusCode, HeaderValue), SkyError> {
	let uri: &str = req.uri().path().as_ref();
	let login = &uri[2..uri.len()];
	let body: String;

	let user = get_user(pool, login.to_string()).await;
	if user.is_ok() {
		body = format!("{}", user?.uuid);
	}
	else {
		body = "Not fond :(".to_owned();
	}

	let restype: HeaderValue = "text/html".parse().unwrap();
	Ok((build_html(&body), StatusCode::OK, restype))
}

fn uri_404() -> (String, StatusCode, HeaderValue) {
	let restype: HeaderValue = "text/html".parse().unwrap();
	(build_html(NF_HTML), StatusCode::NOT_FOUND, restype)
}

async fn init_tables(pool: DBPool) -> Res<(), SkyError> {
	let mut con = pool.get().await.unwrap();

	let _ = con.query_parse::<()>(&query!("CREATE SPACE IF NOT EXISTS bitauth")).await;
	let _ = con.query_parse::<()>(&query!(r#"
		CREATE MODEL IF NOT EXISTS bitauth.users(
			login: string,
			uuid: string,
			password: string,
			email: string,
			tokens: list {type: string}
		)
	"#)).await;
	let _ = con.query_parse::<()>(&query!(r#"
		CREATE MODEL IF NOT EXISTS bitauth.sites(
			uuid: string,
			uid: string,
			domain: string,
			pkey: binary,
			skey: binary
		)
	"#)).await;
	let _ = con.query_parse::<()>(&query!(r#"
		CREATE MODEL IF NOT EXISTS bitauth.tokens(
			uuid: string,
			uid: string,
			sid: string,
			ref: string,
			refend: uint32
		)
	"#)).await;

	let q = con.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = 0")).await;
	if q.is_err() {
		let (skey, pkey) = rsa_gen();
		let _ = con.query_parse::<()>(&query!(
			"INSERT INTO bitauth.sites { uuid: ?, uid: ?, domain: ?, pkey: ?, skey: ? }",
			"0", "0", "", pkey, skey
		)).await;
	};

	Ok(())
}

async fn get_user(pool: DBPool, login: String) -> Res<Users, SkyError> {
	let mut con = pool.get().await.unwrap();

	let q = con
		.query_parse::<Users>(&query!("SELECT * FROM bitauth.users WHERE login = ?", login))
		.await;

	Ok(q?)
}

async fn login_user(pool: DBPool, data: HashMap<String, String>) -> Res<(String, String), SkyError> {
	let mut err = false;
	let mut ret: (String, String) = Default::default();
	let mut con = pool.get().await.unwrap();

	let login = &data.get("login").unwrap().trim().to_lowercase();
	let pass = data.get("password").unwrap().trim();

	let q = con
		.query_parse::<Users>(&query!("SELECT * FROM bitauth.users WHERE login = ?", login))
		.await;
	if q.is_err() { err = true };

	if !err {
		let q = q.unwrap();
		if bcrypt::verify(pass, q.password.as_str()).unwrap() {
			let token = format!("{}", uuid_v4().as_hyphenated());
			let reftoken = format!("{}", uuid_v4().as_hyphenated());
			let time = time();
			let uuid: String = q.uuid;
			let login: String = q.login;

			let _ = con.query_parse::<()>(&query!(
				"INSERT INTO bitauth.tokens { uuid: ?, uid: ?, sid: ?, ref: ?, refend: ? }",
				token.clone(), uuid.clone(), "0", reftoken.clone(), time + REFRESH_LIFETIME
			)).await;
			let _ = con.query_parse::<()>(&query!(
				"UPDATE bitauth.users SET tokens += ? WHERE login = ?",
				token.clone(), login.clone()
			)).await;

			ret = (
				jwt_sign(pool.clone(), json!({
					"login": login.clone(),
					"uuid": uuid.clone(),
					"iat": time,
					"exp": time + TOKEN_LIFETIME
				})).await.unwrap(),
				jwt_sign(pool.clone(), json!({
					"uuid": token.clone(),
					"iat": time,
					"ref": reftoken.clone(),
					"exp": time + REFRESH_LIFETIME
				})).await.unwrap()
			);
		}
	}

	Ok(ret)
}

async fn create_user(pool: DBPool, data: HashMap<String, String>) -> Res<bool, SkyError> {
	let mut ret = true;
	let mut con = pool.get().await.unwrap();

	let login = &data.get("login").unwrap().trim().to_lowercase();
	let email = &data.get("email").unwrap().trim().to_lowercase();
	let pass = data.get("password").unwrap().trim();
	let pass2 = data.get("password2").unwrap().trim();

	if pass != pass2 { ret = false };
	if pass.len() < 8 { ret = false };

	let q = con
		.query_parse::<(String,)>(&query!("SELECT uuid FROM bitauth.users WHERE login = ?", login))
		.await;
	if q.is_ok() { ret = false };

	if ret {
		let uuid = format!("{}", uuid_v4().as_hyphenated());
		let pass = bcrypt::hash(pass, 12).unwrap();

		let q = con.query_parse::<()>(&query!(
			r#"INSERT INTO bitauth.users {
				login: ?,
				uuid: ?,
				password: ?,
				email: ?,
				tokens: []
			}"#,
			login,
			uuid,
			pass,
			email,
		)).await;

		if q.is_err() { ret = false }
	}

	Ok(ret)
}

async fn get_body_from_request(mut req: Request<Incoming>) -> Result<Vec<u8>> {
	let mut body: Vec<u8> = vec![];
	while let Some(next) = req.frame().await {
		let frame = next?;
		if let Some(chunk) = frame.data_ref() {
			body.extend_from_slice(chunk);
		}
	}

	Ok(body)
}

fn double_split(body: String, first: &str, second: &str) -> HashMap<String, String> {
	body.split(first)
		.filter_map(|c| {
			c.split_once(second)
				.map(|(l, r)| (
					l.trim().to_owned(),
					format!("{}", url_decode(r).expect("UTF-8")).trim().to_owned()
				))
		})
		.collect::<HashMap<String, String>>()
}

fn rsa_gen() -> (Vec<u8>, Vec<u8>) {
	let mut rng = rand::thread_rng();
	let bits = 2048;
	let skey = RsaPrivateKey::new(&mut rng, bits).expect("RSA err");
	let pkey = RsaPublicKey::from(&skey);

	(
		skey.to_pkcs1_der().unwrap().as_bytes().to_vec(),
		pkey.to_pkcs1_der().unwrap().as_bytes().to_vec()
	)
}

async fn jwt_sign(pool: DBPool, data: Json) -> Result<String> {
	let mut con = pool.get().await.unwrap();

	let q = con
		.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = ?", "0"))
		.await;
	if q.is_err() { return Ok("".to_owned()) };
	let skey = q.unwrap().skey;

	let skey = EncodingKey::from_rsa_der(&skey);
	let header = Header::new(Algorithm::RS512);
	let token = jwt_encode(&header, &data, &skey)?;

	Ok(token)
}

async fn jwt_verify(pool: DBPool, token: &str) -> Result<TokenData<Json>> {
	let mut con = pool.get().await.unwrap();
	let mut ret = TokenData { header: Header::new(Algorithm::RS512), claims: json![{}] };

	let q = con
		.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = ?", "0"))
		.await;
	if q.is_err() { return Ok(ret) };
	let pkey = q.unwrap().pkey;

	let pkey = DecodingKey::from_rsa_der(&pkey);
	let token = jwt_decode::<Json>(&token, &pkey, &Validation::new(Algorithm::RS512));

	match token.is_ok() {
		true => ret = token.clone().unwrap(),
		_ => {}
	}

	Ok(ret)
}

fn uuid_v4() -> Uuid {
	Uuid::new_v4()
}

fn time() -> u32 {
	SystemTime::now()
		.duration_since(UNIX_EPOCH)
		.unwrap()
		.as_secs() as u32
}

fn time_ns() -> u128 {
	SystemTime::now()
		.duration_since(UNIX_EPOCH)
		.unwrap()
		.as_micros()
}