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

#[derive(Parser)]
pub struct Args {
    #[arg(long)]
    sock: String,
    #[arg(long)]
    pem_dir: PathBuf,
}

#[tokio::main]
async fn main() {
    println!("Starting...");
    let corpus = include_str!("/usr/share/dict/words")
        .split_whitespace()
        .collect::<Vec<_>>();

    let mut freq = HashMap::<String, HashMap<String, usize>>::default();
    include_str!("../wap.txt")
        .split_whitespace()
        .map(|s| s.trim_matches(|c: char| !c.is_alphabetic()))
        .tuple_windows()
        .for_each(|(a, b)| {
            *freq
                .entry(a.to_string())
                .or_default()
                .entry(b.to_string())
                .or_default() += 1
        });
    let mut freq = freq
        .into_iter()
        .map(|(k, v)| (k, v.into_iter().collect::<Vec<_>>()))
        .collect::<HashMap<_, _>>();
    println!("Generated frequencies.");

    let mut counter = Arc::new(AtomicU64::new(0));

    // build our application with a single route
    let app = {
        let counter = counter.clone();
        Router::new().route(
            "/{id}",
            get(|Path(id): Path<String>| async move {
                tokio::time::sleep(Duration::from_millis(random_range(200..1000))).await;

                let mut seed = [0; 32];
                for (i, b) in id
                    .bytes()
                    .chain(core::iter::repeat(0))
                    .take(seed.len())
                    .enumerate()
                {
                    seed[i] = b;
                }
                let mut rng = &mut ChaCha8Rng::from_seed(seed);

                fn choose_word<'a, 'b>(corpus: &'b [&'b str], rng: &'a mut ChaCha8Rng) -> &'b str {
                    corpus.choose(rng).unwrap()
                };

                let title = (0..rng.gen_range(2..10))
                    .map(|_| choose_word(&corpus, rng))
                    .join(" ");

                let mut word = freq.keys().choose(rng).unwrap();
                let mut content = word.clone();
                for _ in 0..rng.gen_range(50..5_000) {
                    // let word = choose_word(&corpus, rng).to_string();
                    word = freq
                        .get(&*word)
                        .and_then(|rhs| rhs.choose_weighted(rng, |(_, n)| *n).ok())
                        .map(|(rhs, _)| rhs)
                        .unwrap_or(word);

                    if rng.gen_bool(0.05) {
                        let url = (0..3).map(|_| choose_word(&corpus, rng)).join("-");
                        content += &format!(" <a href=\"{}\">{}</a>", url, word);
                    } else if rng.gen_bool(0.01) {
                        content += ".<br>";
                    } else {
                        content += " ";
                        content += &word;
                    };
                }
                counter.fetch_add(1, Ordering::Relaxed);
                Html(format!(
                    "<!DOCTYPE html>
                <html>
                <head>
                <title>{title}</title>
                </head>
                <body>

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

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

    // run our app with hyper, listening globally on port 3000
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    println!("Started server.");

    let mut interval = tokio::time::interval(Duration::from_secs(20));
    tokio::spawn(async move {
        loop {
            interval.tick().await;
            println!(
                "Served bollocks to {} clients!",
                counter.load(Ordering::Relaxed)
            );
        }
    });

    let args = Args::parse();

    let config = RustlsConfig::from_pem_file(
        args.pem_dir.clone().join("cert.pem"),
        args.pem_dir.clone().join("key.pem"),
    )
    .await
    .unwrap();
    bind_rustls(args.addr.parse().unwrap(), config)
        .serve(app.into_make_service())
        .await
        .unwrap();
}
