use axum::{
	Router,
	extract::{ConnectInfo, Path, State},
	http::{HeaderMap, HeaderValue, StatusCode, header},
	response::{IntoResponse, Response},
	routing::get,
};
use chrono::{DateTime, TimeDelta, Utc};
use libservice::ServiceConnectInfo;
use lru::LruCache;
use parking_lot::Mutex;
use std::{collections::HashMap, num::NonZero, pin::Pin, sync::Arc, time::Instant};
use tower_http::compression::{CompressionLayer, DefaultPredicate};
use tracing::trace;

use crate::{ClientInfo, RequestContext};

#[derive(Clone)]
pub struct Rendered {
	pub code: StatusCode,
	pub headers: HeaderMap,
	pub body: Vec<u8>,

	pub ttl: Option<TimeDelta>,
}

pub trait Servable: Send + Sync {
	fn render<'a>(
		&'a self,
		ctx: &'a RequestContext,
	) -> Pin<Box<dyn Future<Output = Rendered> + 'a + Send + Sync>>;
}

pub struct PageServer {
	/// If true, expired pages will be rerendered before being sent to the user.
	/// If false, requests never trigger rerenders. We rely on the rerender task.
	///
	/// If true, we deliver fresher pages but delay responses.
	/// TODO: replace this with a smarter rendering strategy?
	never_rerender_on_request: bool,

	/// Map of `{ route: page }`
	pages: Arc<Mutex<HashMap<String, Arc<dyn Servable>>>>,

	/// Map of `{ route: (page data, expire time) }`
	///
	/// We use an LruCache for bounded memory usage.
	page_cache: Mutex<LruCache<RequestContext, (Rendered, DateTime<Utc>)>>,
}

impl PageServer {
	pub fn new() -> Arc<Self> {
		#[expect(clippy::unwrap_used)]
		let cache_size = NonZero::new(128).unwrap();

		Arc::new(Self {
			pages: Arc::new(Mutex::new(HashMap::new())),
			page_cache: Mutex::new(LruCache::new(cache_size)),
			never_rerender_on_request: true,
		})
	}

	pub fn add_page<S: Servable + 'static>(&self, route: impl Into<String>, page: S) -> &Self {
		#[expect(clippy::expect_used)]
		let route = route
			.into()
			.strip_prefix("/")
			.expect("route must start with /")
			.to_owned();

		self.pages.lock().insert(route, Arc::new(page));
		self
	}

	/// Re-render the page at `route`, regardless of cache state.
	/// Does nothing if there is no page at `route`.
	///
	/// Returns the rendered page's content.
	async fn render_page(
		&self,
		reason: &'static str,
		route: &str,
		ctx: RequestContext,
	) -> Option<(Rendered, Option<DateTime<Utc>>)> {
		let now = Utc::now();
		let start = Instant::now();
		let page = match self.pages.lock().get(route) {
			Some(x) => x.clone(),
			None => return None,
		};

		trace!(
			message = "Rendering page",
			route,
			reason,
			lock_time_ms = start.elapsed().as_millis()
		);

		let rendered = page.render(&ctx).await;
		//let html = (self.render_page)(&page, &req_ctx).await.0;

		let mut expires = None;
		if let Some(ttl) = rendered.ttl {
			expires = Some(now + ttl);
			self.page_cache
				.lock()
				.put(ctx, (rendered.clone(), now + ttl));
		}

		let elapsed = start.elapsed().as_millis();
		trace!(message = "Rendered page", route, reason, time_ms = elapsed);
		return Some((rendered, expires));
	}

	async fn handler(
		Path(route): Path<String>,
		State(state): State<Arc<Self>>,
		ConnectInfo(addr): ConnectInfo<ServiceConnectInfo>,
		headers: HeaderMap,
	) -> Response {
		let start = Instant::now();
		let client_info = ClientInfo::from_headers(&headers);
		let ua = headers
			.get("user-agent")
			.and_then(|x| x.to_str().ok())
			.unwrap_or("");

		trace!(
			message = "Serving route",
			route,
			addr = ?addr.addr,
			user_agent = ua,
			device_type = ?client_info.device_type
		);

		// Normalize url with redirect
		if route.ends_with('/') || route.contains("//") || route.starts_with('/') {
			let mut new_route = route.clone();
			while new_route.contains("//") {
				new_route = new_route.replace("//", "/");
			}
			let new_route = new_route.trim_matches('/');

			trace!(
				message = "Redirecting",
				route,
				new_route,
				addr = ?addr.addr,
				user_agent = ua,
				device_type = ?client_info.device_type
			);

			let mut headers = HeaderMap::with_capacity(2);

			let new_route = match HeaderValue::from_str(&format!("/{new_route}")) {
				Ok(x) => x,
				Err(_) => {
					// Be extra careful, this is user-provided data
					return StatusCode::BAD_REQUEST.into_response();
				}
			};

			headers.append(header::LOCATION, new_route);
			headers.append("Accept-CH", HeaderValue::from_static("Sec-CH-UA-Mobile"));
			return (StatusCode::PERMANENT_REDIRECT, headers).into_response();
		}

		let ctx = RequestContext {
			client_info,
			route: format!("/{route}"),
		};

		let now = Utc::now();
		let mut html_expires = None;

		// Get from cache, if available
		if let Some((html, expires)) = state.page_cache.lock().get(&ctx)
			&& (*expires > now || state.never_rerender_on_request)
		{
			html_expires = Some((html.clone(), Some(*expires)));
		};

		if html_expires.is_none() {
			html_expires = match state.render_page("request", &route, ctx).await {
				Some(x) => Some(x.clone()),
				None => {
					trace!(
						message = "Not found",
						route,
						addr = ?addr.addr,
						user_agent = ua,
						device_type = ?client_info.device_type
					);

					trace!(
						message = "Served route",
						route,
						addr = ?addr.addr,
						user_agent = ua,
						device_type = ?client_info.device_type,
						time_ms = start.elapsed().as_millis()
					);

					return StatusCode::NOT_FOUND.into_response();
				}
			};
		}

		#[expect(clippy::unwrap_used)]
		let (mut html, expires) = html_expires.unwrap();

		let max_age = match expires {
			Some(expires) => (expires - now).num_seconds().max(1),
			None => 1,
		};

		#[expect(clippy::unwrap_used)]
		html.headers.insert(
			header::CACHE_CONTROL,
			// immutable;  public/private
			HeaderValue::from_str(&format!("immutable, public, max-age={}", max_age)).unwrap(),
		);

		html.headers
			.insert("Accept-CH", HeaderValue::from_static("Sec-CH-UA-Mobile"));

		trace!(
			message = "Served route",
			route,
			addr = ?addr.addr,
			user_agent = ua,
			device_type = ?client_info.device_type,
			time_ms = start.elapsed().as_millis()
		);

		return (html.code, html.headers, html.body).into_response();
	}

	pub fn into_router(self: Arc<Self>) -> Router<()> {
		let compression: CompressionLayer = CompressionLayer::new()
			.br(true)
			.deflate(true)
			.gzip(true)
			.zstd(true)
			.compress_when(DefaultPredicate::new());

		Router::new()
			.route(
				"/",
				get(|state, conn, headers| async {
					Self::handler(Path(String::new()), state, conn, headers).await
				}),
			)
			.route("/{*path}", get(Self::handler))
			.layer(compression)
			.with_state(self)
	}
}
