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 tower::ServiceBuilder;
use tower_http::services::ServeDir;
use hickory_resolver::Name;
use hickory_resolver::TokioAsyncResolver;

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

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

use lib_humus::TemplateEngineLoader;
use lib_humus::read_toml_from_file;
use lib_humus::HumusEngine;

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

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

type TemplatingEngine = HumusEngine<View,QuerySettings,ResponseFormat>;

#[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<Arc<str>>,
}

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

#[derive(Serialize, 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: Arc<str>,
}

struct ServiceSharedState {
	templating_engine:    TemplatingEngine,
	dns_resolvers:        HashMap<Arc<str>,TokioAsyncResolver>,
	dns_resolver_aliases: HashMap<Arc<str>,Arc<str>>,
	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>,
	default_resolver: Arc<str>,
}

#[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;
}

#[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) {
				Ok(c) => c,
				Err(e) => {
					println!("Could not read confuration file!");
					println!("{e}");
					println!("Exiting ...");
					::std::process::exit(1);
				}
			}
		},
		None => Default::default(),
	};
	
	// Initalize Tera templates
	let template_loader = TemplateEngineLoader::new(
		config.template.template_location.clone(),
		config.template.extra_config.clone()
	)
		.cli_template_location(cli_args.template_location)
		.cli_extra_config_location(cli_args.extra_config);
		

	let templating_engine = match template_loader.load_templates() {
		Ok(t) => t.into(),
		Err(e) => {
			println!("{e}");
			::std::process::exit(1);
		}
	};

	// Static file directory

	let static_file_directory = template_loader.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 resolvers ...");

	let mut dns_resolver_selectables = Vec::<Selectable>::new();
	let mut dns_resolver_map: HashMap<Arc<str>,TokioAsyncResolver> = HashMap::new();
	let mut dns_resolver_aliases: HashMap<Arc<str>,Arc<str>> = HashMap::new();

	if config.dns.enable_system_resolver {
		println!("Initalizing System resolver ...");
		let res = TokioAsyncResolver::tokio_from_system_conf();
		let resolver = match res {
			Ok(resolver) => resolver,
			Err(e) => {
				println!("Error while setting up dns resolver: {e}");
				::std::process::exit(1);
			}
		};

		dns_resolver_map.insert(config.dns.system_resolver_id.clone(), resolver);
		dns_resolver_selectables.push(Selectable {
			id: config.dns.system_resolver_id.clone(),
			name: config.dns.system_resolver_name.clone(),
			weight: config.dns.system_resolver_weight,
		});
	}

	for (key, resolver_config) in &config.dns.resolver {
		println!("Initalizing {} resolver ...", key);
		let resolver = TokioAsyncResolver::tokio(
			resolver_config.to_trust_resolver_config(),
			Default::default()
		);
		dns_resolver_map.insert(key.clone(), resolver);
		dns_resolver_selectables.push(Selectable {
			id: key.clone(),
			name: resolver_config.display_name.clone(),
			weight: resolver_config.weight,
		});
		for alias in &resolver_config.aliases {
			dns_resolver_aliases.insert(alias.clone(),key.clone());
		}
	}
	
	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_resolvers: dns_resolver_map,
		dns_resolver_aliases: dns_resolver_aliases,
		asn_db: asn_db,
		location_db: location_db,
		config: config.clone(),
	});

	dns_resolver_selectables.sort_by(|a,b| b.weight.cmp(&a.weight));
	let default_resolver = dns_resolver_selectables.get(0)
		.map(|s| s.id.clone() )
		.unwrap_or("none".into());
	let derived_config = DerivedConfiguration {
		dns_resolver_selectables: dns_resolver_selectables,
		default_resolver: default_resolver,
	};

	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 = derived_config.default_resolver;

	if let Some(resolver_id) = query.dns {
		dns_resolver_id = resolver_id.into();
	} else if let Some(cookie_header) = cookie_header {
		if let Some(resolver_id) = cookie_header.0.get("dns_resolver") {
			dns_resolver_id = resolver_id.into();
		}
	}
	
	// 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::Text);
					break;
				}
			}
		}
	}
	// Add the request settings extension
	req.extensions_mut().insert(QuerySettings{
		format: format.unwrap_or(ResponseFormat::Html),
		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,
	)
}

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 are an awesome Creature!".to_string()
		},
	)
}


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".into(), &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,
		}
	)
}


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},
			)
		}
	}

	if let Some(via_cap) = VIA_REGEX.captures(&search_query) {
		if let Some(via) = via_cap.get(1) {
			let state = Arc::clone(&arc_state);
			if state.dns_resolvers.contains_key(via.as_str()) {
				settings.dns_resolver_id = via.as_str().into();
			} else if let Some(alias) = state.dns_resolver_aliases.get(via.as_str()) {
				settings.dns_resolver_id = alias.clone();
			}
		}
		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,
	)
}


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.as_str()) {
		state.templating_engine.render_view(
			&settings,
			View::DnsResolver{ config: resolver.clone() },
		)
	} else {
		state.templating_engine.render_view(
			&settings,
			View::NotFound,
		)
	}
}

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}
	)
}

async fn get_ip_result(
	address: &IpAddr,
	lang: &String,
	dns_resolver_name: &Arc<str>,
	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<Arc<str>> = 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}
	)

}

async fn get_dig_result(
	dig_query:         &String,
	dns_resolver_name: &Arc<str>,
	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 let Ok(domain_name) = Name::from_str_relaxed(name.to_owned()+".") {
			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,
							&Name::from_ascii("example.org.").expect("Static Dummy Name"),
							do_full_lookup).await;
				return DigResult {
					records: DnsLookupResult{ nxdomain: true , ..Default::default() },
					idn: idna_name,
					partial_lookup: !do_full_lookup,
					used_dns_resolver: dns_resolver_name.clone(),
				}
			} else {
				return DigResult {
					records: simple_dns::lookup(
						&dns_resolver,
						&domain_name,
						do_full_lookup).await,
					idn: idna_name,
					partial_lookup: !do_full_lookup,
					used_dns_resolver: dns_resolver_name.clone(),
				}
			}
		} else {
			// Invalid domain name
			return DigResult {
				records: DnsLookupResult{
					invalid_name: true,
					.. Default::default()
				},
				idn: idna_name,
				partial_lookup: !do_full_lookup,
				used_dns_resolver: dns_resolver_name.clone(),
			}
		}
	} else {
		// Unknown resolver name
		return DigResult {
			records: DnsLookupResult{
				unkown_resolver: true,
				.. Default::default()
			},
			idn: idna_name,
			partial_lookup: !do_full_lookup,
			used_dns_resolver: "unkown_resolver".into(),
		}
	}
}