use axum::{
	extract::{
		self,
		Query,
		State,
		Extension,
	},
	headers,
	http::Request,
	handler::Handler,
	middleware::{self, Next},
	response::Response,
	Router,
	routing::get,
	TypedHeader,
};
use axum_client_ip::SecureClientIp;
use clap::Parser;
use lazy_static::lazy_static;
use regex::Regex;
use serde::{Deserialize,Serialize};
use tera::Tera;
use tower::ServiceBuilder;
use tower_http::services::ServeDir;
use trust_dns_resolver::TokioAsyncResolver;

use tokio::signal::unix::{
	signal,
	SignalKind,
};
use tokio::task;

use std::collections::HashMap;
use std::fs;
use std::net::IpAddr;
use std::sync::Arc;

mod config;
mod geoip;
mod ipinfo;
mod ratelimit;
mod settings;
mod simple_dns;
mod templating_engine;
mod idna;

use crate::geoip::{
	QueryAsn,
	QueryLocation,
	AsnResult,
	LocationResult,
};
use crate::idna::IdnaName;
use crate::simple_dns::DnsLookupResult;
use crate::settings::*;
use crate::templating_engine::View;
use crate::ipinfo::{AddressCast,AddressInfo,AddressScope};

#[derive(Deserialize, Serialize, Clone)]
pub struct SettingsQuery {
	format: Option<ResponseFormat>,
	lang:   Option<String>,
	dns:	Option<String>,
}

#[derive(Deserialize, Serialize, Clone)]
pub struct SearchQuery {
	query: Option<String>,
}

#[derive(Serialize, Clone)]
pub struct IpResult {
	address:  IpAddr,
	hostname: Option<String>,
	asn:      Option<AsnResult>,
	location: Option<LocationResult>,
	ip_info:  AddressInfo,
	used_dns_resolver: Option<String>,
}

// We need this one to hide the partial lookup field when irelevant
pub fn not(b: &bool) -> bool { !b }

#[derive(Serialize, Default, Clone)]
pub struct DigResult {
	records: simple_dns::DnsLookupResult,
	#[serde(skip_serializing_if = "IdnaName::was_ascii")]
	idn: IdnaName,
	#[serde(skip_serializing_if = "not")]
	partial_lookup: bool,
	used_dns_resolver: String,
}

struct ServiceSharedState {
	templating_engine: templating_engine::Engine,
	dns_resolvers:     HashMap<String,TokioAsyncResolver>,
	asn_db:            geoip::MMDBCarrier,
	location_db:       geoip::MMDBCarrier,
	config:            config::EchoIpServiceConfig,
}

// Stores configuration that is derived from the original configuration
#[derive(Clone)]
struct DerivedConfiguration {
	dns_resolver_selectables: Vec<Selectable>,
}

#[derive(Parser)]
#[command(author, version, long_about="A web service that tells you your ip-address and more …")]
struct CliArgs {
	#[arg(short, long)]
	config: Option<String>,
	#[arg(short, long)]
	listen_on: Option<String>,
	#[arg(short, long)]
	template_location: Option<String>,
	#[arg(short,long)]
	extra_config: Option<String>,
	#[arg(short,long)]
	static_location: Option<String>,
}

fn match_domain_hidden_list(domain: &String, hidden_list: &Vec<String>) -> bool {
	let name = domain.trim_end_matches(".");
	for suffix in hidden_list {
		if name.ends_with(suffix) {
			return true;
		}
	}
	return false;
}

fn read_toml_from_file<T: for<'de> serde::Deserialize<'de>>(path: &String) -> Option<T> {
	let text = match fs::read_to_string(path) {
		Ok(t) => t,
		Err(e) => {
			println!("Error while reading file '{path}': {e}");
			return None;
		}
	};
	match toml::from_str(&text) {
		Ok(t) => Some(t),
		Err(e) => {
			println!("Unable to parse file '{path}':\n{e}");
			return None;
		}
	}
}

#[tokio::main]
async fn main() {
	// Parse Command line arguments
	let cli_args = CliArgs::parse();

	// Read configuration file
	let config: config::EchoIpServiceConfig = match cli_args.config {
		Some(config_path) => {
			match read_toml_from_file::<config::EchoIpServiceConfig>(&config_path) {
				Some(c) => c,
				None => {
					println!("Could not read confuration file, exiting.");
					::std::process::exit(1);
				}
			}
		},
		None => Default::default(),
	};
	
	// Initalize Tera templates
	let mut template_base_dir = match cli_args.template_location {
		Some(template_base_dir) => template_base_dir,
		None => (&config.template.template_location).to_owned(),
	};
	if !template_base_dir.ends_with("/") {
		template_base_dir = template_base_dir + "/";
	}
	let template_extra_config = match &cli_args.extra_config {
		Some(path) => read_toml_from_file(path),
		None => match &config.template.extra_config {
			Some(path) => read_toml_from_file(path),
			None => {
				println!("Trying to read default template configuration ...");
				println!("(If this fails that may be ok, depending on your template)");
				read_toml_from_file(&(template_base_dir.clone()+"extra.toml"))
			},
		},
	};
	let template_glob = template_base_dir.clone()+"*";
	println!("Parsing Templates from '{}' ...", &template_glob);
	let res = Tera::new((template_glob).as_str());
	let tera = match res {
		Ok(t) => t,
		Err(e) => {
			println!("Template parsing error(s): {}", e);
			::std::process::exit(1);
		}
	};

	let templating_engine = templating_engine::Engine{
		tera: tera,
		template_config: template_extra_config,
	};

	// Static file directory

	let static_file_directory = cli_args.static_location.unwrap_or(
		config.server.static_location.clone().unwrap_or(
			template_base_dir+"/static"
		)
	);

	println!("Static files will be served from: {static_file_directory}");

	// Initalize GeoIP Database

	let asn_db = geoip::MMDBCarrier::new(
		"GeoIP ASN Database".to_string(),
		config.geoip.asn_database.clone()
	);

	asn_db.reload_database().ok();

	let location_db = geoip::MMDBCarrier::new(
		"GeoIP Location Database".to_string(),
		config.geoip.location_database.clone()
	);

	location_db.reload_database().ok();
	
	// Initalize DNS resolver with os defaults
	println!("Initalizing dns resolver ...");

	let mut dns_resolver_selectables = Vec::<Selectable>::new();

	println!("Initalizing System resolver ...");
	let res = TokioAsyncResolver::tokio_from_system_conf();
	//let res = TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
	let dns_resolver = match res {
		Ok(resolver) => resolver,
		Err(e) => {
			println!("Error while setting up dns resolver: {e}");
			::std::process::exit(1);
		}
	};

	dns_resolver_selectables.push(Selectable {
		id: "default".to_string(),
		name: "System".to_string()
	});

	//FIXME: Not release ready,must be configurable and have better error handling.
	println!("Initalizing Quad9 resolver ...");
	let quad9_resolver = TokioAsyncResolver::tokio(
		trust_dns_resolver::config::ResolverConfig::quad9_tls(),
		Default::default()
	).unwrap();
	dns_resolver_selectables.push(Selectable {
		id: "quad9".to_string(),
		name: "Quad9".to_string()
	});
	println!("Initalizing Google resolver ...");
	let google_resolver = TokioAsyncResolver::tokio(
		trust_dns_resolver::config::ResolverConfig::google(),
		Default::default()
	).unwrap();
	dns_resolver_selectables.push(Selectable {
		id: "google".to_string(),
		name: "Google".to_string()
	});
	println!("Initalizing Cloudflare resolver ...");
	let cloudflare_resolver = TokioAsyncResolver::tokio(
		trust_dns_resolver::config::ResolverConfig::cloudflare_tls(),
		Default::default()
	).unwrap();
	dns_resolver_selectables.push(Selectable {
		id: "cloudflare".to_string(),
		name: "Cloudflare".to_string()
	});

	let mut dns_resolver_map: HashMap<String,TokioAsyncResolver> = HashMap::new();

	for (key, resolver_config) in &config.dns.resolver {
		println!("Initalizing {} resolver ...", key);
		let resolver = TokioAsyncResolver::tokio(
			resolver_config.to_trust_resolver_config(),
			Default::default()
		).unwrap();
		dns_resolver_map.insert(key.clone(), resolver);
		dns_resolver_selectables.push(Selectable {
			id: key.clone(),
			name: resolver_config.display_name.clone(),
		});
	}
	
	dns_resolver_map.insert("default".to_string(), dns_resolver);
	dns_resolver_map.insert("quad9".to_string(), quad9_resolver);
	dns_resolver_map.insert("google".to_string(), google_resolver);
	dns_resolver_map.insert("cloudflare".to_string(), cloudflare_resolver);

	let listen_on = config.server.listen_on;
	let ip_header = config.server.ip_header.clone();

	// Initialize shared state
	let shared_state = Arc::new(
		ServiceSharedState {
		templating_engine: templating_engine,
		//dns_resolver: dns_resolver,
		dns_resolvers: dns_resolver_map,
		asn_db: asn_db,
		location_db: location_db,
		config: config.clone(),
	});

	let derived_config = DerivedConfiguration {
		dns_resolver_selectables: dns_resolver_selectables,
	};

	let signal_usr1_handlers_state = shared_state.clone();

	task::spawn(async move {
		println!("Trying to register USR1 signal for reloading geoip databases");
		let mut signal_stream = match signal(SignalKind::user_defined1()) {
			Ok(signal_stream) => signal_stream,
			Err(e) => {
				println!("Error while registring signal handler: {e}");
				println!("Continuing without ...");
				return;
			}
		};
		loop {
			if None == signal_stream.recv().await { return; }
			println!("Received signal USR1, reloading geoip databses!");
			signal_usr1_handlers_state.location_db.reload_database().ok();
			signal_usr1_handlers_state.asn_db.reload_database().ok();
		}
	});

	// Initalize axum server
    let app = Router::new()
	    .route("/", get(handle_default_route))
	    .route("/dig/:name", get(handle_dig_route_with_path))
	    .route("/ip/:address", get(handle_ip_route_with_path))
	    .route("/dns_resolver/:resolver", get(handle_dns_resolver_route_with_path))
	    .route("/dns_resolver", get(handle_dns_resolver_route))
	    .route("/ua", get(user_agent_handler))
    	.route("/hi", get(hello_world_handler))
    	.fallback_service(
	    	ServeDir::new(static_file_directory)
		    	.fallback(not_found_handler.with_state(shared_state.clone()))
    	)
		.with_state(shared_state)
    	.layer(
			ServiceBuilder::new()
				.layer(ip_header.into_extension())
				.layer(ratelimit::build_rate_limiting_state(
					config.ratelimit.per_minute, config.ratelimit.burst))
				.layer(middleware::from_fn(ratelimit::rate_limit_middleware))
				.layer(Extension(config))
				.layer(Extension(derived_config))
				.layer(middleware::from_fn(settings_query_middleware))
    	)
    	;

	println!("Starting Server on {} ...",listen_on);

    axum::Server::bind(&listen_on)
        .serve(app.into_make_service_with_connect_info::<std::net::SocketAddr>())
        .await
        .unwrap();
}


async fn settings_query_middleware<B>(
	Query(query): Query<SettingsQuery>,
	Extension(config): Extension<config::EchoIpServiceConfig>,
	Extension(derived_config): Extension<DerivedConfiguration>,
	cookie_header: Option<TypedHeader<headers::Cookie>>,
	user_agent_header: Option<TypedHeader<headers::UserAgent>>,
	mut req: Request<B>,
	next: Next<B>
) -> Response {
	let mut format = query.format;
	let mut dns_resolver_id = config.dns.default_resolver;

	if let Some(resolver_id) = query.dns {
		dns_resolver_id = resolver_id;
	} else if let Some(cookie_header) = cookie_header {
		if let Some(resolver_id) = cookie_header.0.get("dns_resolver") {
			dns_resolver_id = resolver_id.to_string();
		}
	}
	
	// Try to guess type from user agent
	if format.is_none() {
		if let Some(TypedHeader(user_agent)) = user_agent_header {
			let ua = user_agent.as_str();
			for tua in config.template.text_user_agents {
				if ua.starts_with(&tua) {
					format = Some(ResponseFormat::TextPlain);
					break;
				}
			}
		}
	}
	// Add the request settings extension
	req.extensions_mut().insert(QuerySettings{
		format: format.unwrap_or(ResponseFormat::TextHtml),
		lang: query.lang.unwrap_or("en".to_string()),
		available_dns_resolvers: derived_config.dns_resolver_selectables,
		dns_resolver_id: dns_resolver_id,
	});
	next.run(req).await
}

async fn not_found_handler(
	State(arc_state): State<Arc<ServiceSharedState>>,
	Extension(settings): Extension<QuerySettings>,
) -> Response {
	let state = Arc::clone(&arc_state);

	state.templating_engine.render_view(
		&settings,
		&View::NotFound,
	).await
}

async fn hello_world_handler(
	State(arc_state): State<Arc<ServiceSharedState>>,
	Extension(settings): Extension<QuerySettings>,
) -> Response {
	let state = Arc::clone(&arc_state);

	state.templating_engine.render_view(
		&settings,
		&View::Message{
			title: "Hey There!".to_string(),
			message: "You,You are an awesome Creature!".to_string()
		},
	).await
}


async fn user_agent_handler(
	TypedHeader(user_agent): TypedHeader<headers::UserAgent>,
) -> String {
	user_agent.to_string()
}

async fn handle_default_route(
	Query(search_query): Query<SearchQuery>,
	State(arc_state): State<Arc<ServiceSharedState>>,
	Extension(settings): Extension<QuerySettings>,
	user_agent_header: Option<TypedHeader<headers::UserAgent>>,
	SecureClientIp(address): SecureClientIp
) -> Response {
	
	let state = Arc::clone(&arc_state);

	if let Some(search_query) = search_query.query {
		if search_query.trim() != "" {
			return handle_search_request(
				search_query,
				false,
				settings,
				state
			).await;
		}
	}

	let result = get_ip_result(&address, &settings.lang, &"default".to_string(), &state).await;

	let user_agent: Option<String> = match user_agent_header {
		Some(TypedHeader(user_agent)) => Some(user_agent.to_string()),
		None => None,
	};

	state.templating_engine.render_view(
		&settings,
		&View::Index{
			result: result,
			user_agent: user_agent,
		}
	).await
}


async fn handle_search_request(
	search_query: String,
	this_should_have_been_an_ip: bool,
	settings: QuerySettings,
	arc_state: Arc<ServiceSharedState>,
) -> Response {

	let mut search_query = search_query.trim().to_string();
	let mut settings = settings;

	lazy_static!{
		static ref ASN_REGEX: Regex = Regex::new(r"^[Aa][Ss][Nn]?\s*(\d{1,7})$").unwrap();
		static ref VIA_REGEX: Regex = Regex::new(r"[Vv][Ii][Aa]\s+(\S+)").unwrap();
	}

	//If someone asked for an asn, give an asn answer
	if let Some(asn_cap) = ASN_REGEX.captures(&search_query) {
		if let Some(asn) = asn_cap.get(1).map_or(None, |m| m.as_str().parse::<u32>().ok()) {
			// Render a dummy template that can at least link to other pages
			let state = Arc::clone(&arc_state);
			return state.templating_engine.render_view(
				&settings,
				&View::Asn{asn: asn},
			).await
		}
	}

	if let Some(via_cap) = VIA_REGEX.captures(&search_query) {
		if let Some(via) = via_cap.get(1).map(|c| c.as_str().to_string()) {
			let state = Arc::clone(&arc_state);
			if state.dns_resolvers.contains_key(&via) {
				settings.dns_resolver_id = via;
			}
		}
		search_query = VIA_REGEX.replace(&search_query,"").trim().to_string();
	}

	// Try to interpret as an IP-Address
	if let Ok(address) = search_query.parse() {
		return handle_ip_request(address, settings, arc_state).await;
	}

	// Fall back to treating it as a hostname
	return handle_dig_request(
		search_query,
		settings,
		arc_state,
		!this_should_have_been_an_ip,
	).await

}

async fn handle_dns_resolver_route(
	Extension(settings): Extension<QuerySettings>,
	State(arc_state): State<Arc<ServiceSharedState>>,
) -> Response {
	let state = Arc::clone(&arc_state);
	state.templating_engine.render_view(
		&settings,
		&View::DnsResolverList,
	).await	
}


async fn handle_dns_resolver_route_with_path(
	Extension(settings): Extension<QuerySettings>,
	State(arc_state): State<Arc<ServiceSharedState>>,
	extract::Path(query): extract::Path<String>,
) -> Response {
	let state = Arc::clone(&arc_state);
	if let Some(resolver) = state.config.dns.resolver.get(&query) {
		state.templating_engine.render_view(
			&settings,
			&View::DnsResolver{ config: resolver.clone() },
		).await
	} else {
		state.templating_engine.render_view(
			&settings,
			&View::NotFound,
		).await
	}
}

async fn handle_ip_route_with_path(
	Extension(settings): Extension<QuerySettings>,
	State(arc_state): State<Arc<ServiceSharedState>>,
	extract::Path(query): extract::Path<String>,
) -> Response {
	if let Ok(address) = query.parse() {
		return handle_ip_request(address, settings, arc_state).await
	} else {
		return handle_search_request(query, true, settings, arc_state).await;
	}
}

async fn handle_ip_request(
	address: IpAddr,
	settings: QuerySettings,
	arc_state: Arc<ServiceSharedState>,
) -> Response {

	let state = Arc::clone(&arc_state);
	let result = get_ip_result(
		&address,
		&settings.lang,
		&settings.dns_resolver_id,
		&state).await;

	state.templating_engine.render_view(
		&settings,
		&View::Ip{result: result}
	).await
}

async fn get_ip_result(
	address: &IpAddr,
	lang: &String,
	dns_resolver_name: &String,
	state: &ServiceSharedState,
) -> IpResult {

	let ip_info = AddressInfo::new(&address);

	if !(ip_info.scope == AddressScope::Global || ip_info.scope == AddressScope::Shared) || ip_info.cast != AddressCast::Unicast {
		if !((ip_info.scope == AddressScope::Private || ip_info.scope == AddressScope::LinkLocal) && state.config.server.allow_private_ip_lookup) {
			return IpResult {
				address: *address,
				hostname: None,
				asn: None,
				location: None,
				ip_info: ip_info,
				used_dns_resolver: None,
			}
		}
	} 

	// do reverse lookup
	let mut hostname: Option<String> = None;
	let mut used_dns_resolver: Option<String> = None;
	if state.config.dns.allow_reverse_lookup {
		if let Some(dns_resolver) = &state.dns_resolvers.get(dns_resolver_name) {
			hostname = simple_dns::reverse_lookup(&dns_resolver, &address).await;
			used_dns_resolver = Some(dns_resolver_name.clone());
		}
	}

	// asn lookup
	let asn_result = state.asn_db.query_asn_for_ip(address);

	// location lookup
	let location_result = state.location_db.query_location_for_ip(
		address,
		&vec![lang, &"en".to_string()]
	);

	// filter reverse lookup
	if let Some(name) = &hostname {
		if match_domain_hidden_list(&name, &state.config.dns.hidden_suffixes) {
			hostname = None;
			used_dns_resolver = None;
		}
	}

	IpResult{
		address: *address,
		hostname: hostname,
		asn: asn_result,
		location: location_result,
		ip_info: ip_info,
		used_dns_resolver: used_dns_resolver,
	}
}

async fn handle_dig_route_with_path(
	Extension(settings): Extension<QuerySettings>,
	State(arc_state): State<Arc<ServiceSharedState>>,
	extract::Path(name): extract::Path<String>,
) -> Response {
	return handle_dig_request(name, settings, arc_state, true).await
}

async fn handle_dig_request(
	dig_query: String,
	settings: QuerySettings,
	arc_state: Arc<ServiceSharedState>,
	do_full_lookup: bool,
) -> Response {

	let state = Arc::clone(&arc_state);

	let dig_result = get_dig_result(
		&dig_query,
		&settings.dns_resolver_id,
		&state,
		do_full_lookup
	).await;

	state.templating_engine.render_view(
		&settings,
		&View::Dig{ query: dig_query, result: dig_result}
	).await

}

async fn get_dig_result(
	dig_query:         &String,
	dns_resolver_name: &String,
	state:             &ServiceSharedState,
	do_full_lookup:     bool,
) -> DigResult {
	let name = &dig_query.trim().trim_end_matches(".").to_string();
	let idna_name = IdnaName::from_string(&name);
	if let Some(dns_resolver) = &state.dns_resolvers.get(dns_resolver_name) {
		if match_domain_hidden_list(&name, &state.config.dns.hidden_suffixes) {
			// Try to hide the fact that we didn't do dns resolution at all
			// We resolve example.org as basic avoidance of timing sidechannels.
			// WARNING: this timing sidechannel avoidance is very crude.
			simple_dns::lookup(
						&dns_resolver,
						&("example.org.".to_string()),
						do_full_lookup).await;
			DigResult {
				records: DnsLookupResult{ nxdomain: true , ..Default::default() },
				idn: idna_name,
				partial_lookup: !do_full_lookup,
				used_dns_resolver: dns_resolver_name.clone(),
			}
		} else {
			DigResult {
				records: simple_dns::lookup(
					&dns_resolver,
					&(idna_name.idn.clone().unwrap_or(name.to_owned())+"."),
					do_full_lookup).await,
				idn: idna_name,
				partial_lookup: !do_full_lookup,
				used_dns_resolver: dns_resolver_name.clone(),
			}
		}
	} else {
		return Default::default();
	}
}