mod generator;

use axum::{extract::Path, response::Html, routing::get, Router};
use axum_server::{bind_rustls, tls_rustls::RustlsConfig};
use clap::Parser;
use hashbrown::HashMap;
use itertools::Itertools;
use rand::{prelude::*, Rng as _};
use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng};
use std::{
    borrow::Cow,
    path::PathBuf,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    time::Duration,
};

pub type Rng = ChaCha8Rng;

#[derive(Parser)]
pub struct Args {
    /// Socket to bind to, defaults to 0.0.0.0:3000
    #[arg(long)]
    sock: Option<String>,
    #[arg(long)]
    /// Path of the certificate .pem
    cert: Option<PathBuf>,
    /// Path of the private key .pem
    #[arg(long)]
    key: Option<PathBuf>,
}

#[inline]
fn create_rng(seed_bytes: impl IntoIterator<Item = u8>) -> Rng {
    let mut seed = [0; 32];
    for (i, b) in seed_bytes.into_iter().take(seed.len()).enumerate() {
        seed[i] = b;
    }
    Rng::from_seed(seed)
}

const COUNT_FILE: &str = "count.txt";

#[tokio::main]
async fn main() {
    let args = Args::parse();

    let generators: Vec<Arc<dyn generator::Generator>> = vec![
        Arc::new(generator::Markov::new(include_str!("../wap.txt"))),
        Arc::new(generator::Ast::new()),
    ];

    let counter = Arc::new(AtomicU64::new(
        if let Some(prev_count) = std::fs::read_to_string(COUNT_FILE)
            .ok()
            .and_then(|s| s.parse().ok())
        {
            prev_count
        } else {
            0
        },
    ));

    let app = {
        let counter = counter.clone();
        Router::new().route(
            "/{id}",
            get(|Path(id): Path<String>| async move {
                // Create a RNG for this path (deterministic, to simulate static pages)
                let mut rng = create_rng(id.bytes());

                // Count the request. Also doubles as the non-deterministic seed
                let count = counter.fetch_add(1, Ordering::Relaxed);

                // Create a RNG for this session (non-deterministic)
                let mut session_rng = create_rng(count.to_le_bytes());

                // Artificially slow down connections as rudimentary DDoS protection, and to use up client resources
                tokio::time::sleep(Duration::from_millis(session_rng.random_range(200..1000)))
                    .await;

                // Choose a bullshit generator from our collection for this page
                let generator = generators.choose(&mut rng).unwrap();

                let title = generator
                    .word_stream(rng.random_range(2..10), &mut rng.clone())
                    .join(" ");

                let stats = format!("Served rubbish to {count} clients so far");

                let content = generator
                    .word_stream(rng.random_range(50..5_000), &mut rng.clone())
                    .fold(String::new(), |mut content, word| {
                        // Small chance of every word becoming a link back into the void
                        if rng.random_bool(0.05) {
                            let url = generator.word_stream(3, &mut rng.clone()).join("-");
                            content += &format!(" <a href=\"{}\">{}</a>", url, word);
                        } else {
                            // Also, a chance for every word to end with a newline. This should probably be controlled by the generator.
                            content += if rng.random_bool(0.01) { ".<br>" } else { " " };
                            content += &word
                        }
                        content
                    });

                Html(format!(
                    "<!DOCTYPE html>
                    <html>
                    <head>
                    <title>{title}</title>
                    </head>
                    <body>

                    <h1>{title}</h1>
                    <p>{stats}</p>
                    <p>{content}</p>

                    </body>
                    </html>"
                ))
            }),
        )
    };

    let mut interval = tokio::time::interval(Duration::from_secs(20));
    tokio::spawn(async move {
        let mut last = 0;
        loop {
            interval.tick().await;

            let count = counter.load(Ordering::Relaxed);

            if count != last {
                last = count;
                println!(
                    "{} Served bollocks to {} clients!",
                    chrono::offset::Local::now(),
                    count,
                );
                let _ = std::fs::write(COUNT_FILE, &format!("{count}"));
            }
        }
    });

    println!("Starting...");
    let sock = args
        .sock
        .as_deref()
        .unwrap_or("0.0.0.0:3000")
        .parse()
        .unwrap();
    if let (Some(cert), Some(key)) = (args.cert, args.key) {
        println!("Enabling TLS...");
        let config = RustlsConfig::from_pem_file(cert, key).await.unwrap();
        bind_rustls(sock, config)
            .serve(app.into_make_service())
            .await
            .unwrap();
    } else {
        println!("WARNING: TLS disabled.");
        axum_server::bind(sock)
            .serve(app.into_make_service())
            .await
            .unwrap()
    }
}
