use chrono::{prelude::*, Utc};
use clap::Parser;
use futures::future;
use futures::StreamExt;
use hyper::{
    server::conn::AddrStream,
    service::{make_service_fn, service_fn},
    Body, Request, Response, Server,
};
use lox_library::bridge_table::{BridgeLine, MAX_BRIDGES_PER_BUCKET};
use lox_library::{BridgeAuth, BridgeDb};

use rdsys_backend::{proto::ResourceDiff, start_stream};
use serde::Deserialize;
use std::env;
use std::{
    convert::Infallible,
    error::Error,
    fs::File,
    io::BufReader,
    net::SocketAddr,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
    time::Duration,
};

mod lox_context;
mod request_handler;
use request_handler::handle;
mod resource_parser;
use resource_parser::parse_resource;

use tokio::{
    signal, spawn,
    sync::{broadcast, mpsc, oneshot},
    time::sleep,
};

async fn shutdown_signal() {
    tokio::signal::ctrl_c()
        .await
        .expect("failed to listen for ctrl+c signal");
    println!("Shut down Lox Server");
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Name/path of the configuration file
    #[arg(short, long, default_value = "config.json")]
    config: PathBuf,

    /// Optional name/path of the db file
    #[arg(short, long)]
    db: Option<PathBuf>,
}

#[derive(Debug, Deserialize)]
struct ResourceInfo {
    endpoint: String,
    name: String,
    token: String,
    types: Vec<String>,
}
// Populate Bridgedb from rdsys

// Rdsys sender creates a ResourceStream with the api_endpoint, resource token and type specified
// in the config.json file.
// TODO: ensure this stream gracefully shutdowns on the ctrl_c command.
async fn rdsys_stream(
    rtype: ResourceInfo,
    tx: mpsc::Sender<ResourceDiff>,
    mut kill: broadcast::Receiver<()>,
) {
    let mut rstream = start_stream(rtype.endpoint, rtype.name, rtype.token, rtype.types)
        .await
        .expect("rdsys stream initialization failed. Start rdsys or check config.json");
    loop {
        tokio::select! {
        res = rstream.next() => {
            match res {
                Some(diff) => tx.send(diff).await.unwrap(),
                None => return,
            }
        },
                _ = kill.recv() => {println!("Shut down rdsys stream"); return},

        }
    }
}

async fn rdsys_bridge_parser(
    rdsys_tx: mpsc::Sender<Command>,
    rx: mpsc::Receiver<ResourceDiff>,
    mut kill: broadcast::Receiver<()>,
) {
    tokio::select! {
        start_bridge_parser = parse_bridges(rdsys_tx, rx) => start_bridge_parser ,
        _ = kill.recv() => {println!("Shut down bridge_parser");},
    }
}

// Parse Bridges receives a ResourceDiff from rdsys_sender and sends it to the
// Context Manager to be parsed and added to the BridgeDB
async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, mut rx: mpsc::Receiver<ResourceDiff>) {
    loop {
        let resourcediff = rx.recv().await.unwrap();
        let cmd = Command::Rdsys { resourcediff };
        rdsys_tx.send(cmd).await.unwrap();
        sleep(Duration::from_secs(1)).await;
    }
}

async fn create_context_manager(
    db_path: Option<PathBuf>,
    context_rx: mpsc::Receiver<Command>,
    mut kill: broadcast::Receiver<()>,
) {
    tokio::select! {
        create_context = context_manager(db_path, context_rx) => create_context,
        _ = kill.recv() => {println!("Shut down context_manager");},
    }
}

fn read_context_from_file<P: AsRef<Path>>(
    path: P,
) -> Result<lox_context::LoxServerContext, Box<dyn Error>> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);
    let context = serde_json::from_reader(reader)?;
    Ok(context)
}

// Context Manager handles the Lox BridgeDB and Bridge Authority, ensuring
// that the DB can be updated from the rdsys stream and client requests
// can be responded to with an updated BridgeDB state
async fn context_manager(db_path: Option<PathBuf>, mut context_rx: mpsc::Receiver<Command>) {
    let context: lox_context::LoxServerContext;
    if let Some(existing_db) = db_path.as_deref() {
        context = read_context_from_file(existing_db).unwrap();
    } else {
        let current_path = env::current_dir().expect("Unable to access current dir");
        if let Some(last_modified_file) = std::fs::read_dir(current_path)
            .expect("Couldn't read local directory")
            .flatten() // Remove failed
            .filter(|f| {
                f.metadata().unwrap().is_file()
                    && (f.file_name().into_string().unwrap().contains("_lox.json"))
            }) // Filter out directories (only consider files)
            .max_by_key(|x| x.metadata().unwrap().modified().unwrap())
        // Get the most recently modified file
        {
            println!("Reading from file {:?}", last_modified_file);
            context = read_context_from_file(&last_modified_file.path()).unwrap();
        } else {
            let new_db = BridgeDb::new();
            let new_ba = BridgeAuth::new(new_db.pubkey);
            context = lox_context::LoxServerContext {
                db: Arc::new(Mutex::new(new_db)),
                ba: Arc::new(Mutex::new(new_ba)),
                extra_bridges: Arc::new(Mutex::new(Vec::new())),
                to_be_replaced_bridges: Arc::new(Mutex::new(Vec::new())),
            };
        }
    }

    while let Some(cmd) = context_rx.recv().await {
        use Command::*;
        match cmd {
            Rdsys { resourcediff } => {
                if let Some(new_resources) = resourcediff.new {
                    let mut count = 0;
                    let mut bucket = [BridgeLine::default(); MAX_BRIDGES_PER_BUCKET];
                    for pt in new_resources {
                        println!("A NEW RESOURCE: {:?}", pt);
                        if let Some(resources) = pt.1 {
                            for resource in resources {
                                let bridgeline = parse_resource(resource);
                                println!("Now it's a bridgeline: {:?}", bridgeline);
                                if context.to_be_replaced_bridges.lock().unwrap().len() > 0 {
                                    println!("BridgeLine to be replaced: {:?}", bridgeline);
                                    let res = context.replace_with_new(bridgeline);
                                    if res == lox_library::ReplaceSuccess::NotFound {
                                        println!(
                                        "BridgeLine not found in bridge_table, already updated {:?}",
                                        bridgeline
                                    );
                                    } else if res == lox_library::ReplaceSuccess::Replaced {
                                        println!(
                                            "BridgeLine successfully replaced: {:?}",
                                            bridgeline
                                        );
                                    } else {
                                        assert!(
                                            res == lox_library::ReplaceSuccess::NotReplaced,
                                            "ReplaceSuccess incorrectly set somehow"
                                        );
                                        // Add the bridge to the list of to_be_replaced bridges in the Lox context and try
                                        // again to replace at the next update (nothing changes in the Lox Authority)
                                        println!("'Gone' BridgeLine NOT replaced, saved for next update! : {:?}", bridgeline);
                                        context.new_to_be_replaced_bridge(bridgeline);
                                    }
                                } else if count < MAX_BRIDGES_PER_BUCKET {
                                    bucket[count] = bridgeline;
                                    count += 1;
                                } else {
                                    // TODO: Decide the circumstances under which a bridge is allocated to an open_inv or spare bucket,
                                    // eventually also do some more fancy grouping of new resources, i.e., by type or region
                                    context.add_openinv_bucket(bucket);
                                    count = 0;
                                    bucket = [BridgeLine::default(); MAX_BRIDGES_PER_BUCKET];
                                }
                            }
                        }
                    }
                    // Handle the extra buckets that were not allocated already
                    if count != 0 {
                        for val in 0..count {
                            if context.extra_bridges.lock().unwrap().len()
                                < (MAX_BRIDGES_PER_BUCKET)
                            {
                                context.append_extra_bridges(bucket[val]);
                            } else {
                                bucket = context.remove_extra_bridges();
                                context.add_spare_bucket(bucket);
                            }
                        }
                    }
                }
                if let Some(changed_resources) = resourcediff.changed {
                    for pt in changed_resources {
                        println!("A NEW CHANGED RESOURCE: {:?}", pt);
                        if let Some(resources) = pt.1 {
                            for resource in resources {
                                let bridgeline = parse_resource(resource);
                                println!("BridgeLine to be changed: {:?}", bridgeline);
                                let res = context.update_bridge(bridgeline);
                                if res {
                                    println!("BridgeLine successfully updated: {:?}", bridgeline);
                                } else {
                                    println!("BridgeLine: {:?} not found in Lox's Bridgetable. Save it as a new resource for now!", bridgeline);
                                    if context.extra_bridges.lock().unwrap().len() < 2 {
                                        context.append_extra_bridges(bridgeline);
                                    } else {
                                        let bucket = context.remove_extra_bridges();
                                        context.add_spare_bucket(bucket);
                                    }
                                }
                            }
                        }
                    }
                }
                // gone resources are not the same as blocked resources.
                // Instead, these are bridges which have either failed to pass tests for some period
                // or have expired bridge descriptors. In both cases, the bridge is unusable, but this
                // is not likely due to censorship. Therefore, we replace gone resources with new resources
                // TODO: create a notion of blocked resources from information collected through various means:
                // https://gitlab.torproject.org/tpo/anti-censorship/censorship-analysis/-/issues/40035
                if let Some(gone_resources) = resourcediff.gone {
                    for pt in gone_resources {
                        println!("A NEW GONE RESOURCE: {:?}", pt);
                        if let Some(resources) = pt.1 {
                            for resource in resources {
                                // If resource last passed tests 3 hours ago, it should be replaced with a working
                                // resource and be removed from the bridgetable. If it has been gone for more than 7 hours,
                                // we should stop trying to remove it from the bridge table and assume it has successfully been
                                // removed already
                                if resource.last_passed < (Utc::now() - chrono::Duration::hours(3))
                                    || resource.last_passed
                                        > (Utc::now() - chrono::Duration::hours(7))
                                {
                                    let bridgeline = parse_resource(resource);
                                    println!("BridgeLine to be replaced: {:?}", bridgeline);
                                    let res = context.replace_with_new(bridgeline);
                                    if res == lox_library::ReplaceSuccess::Replaced {
                                        println!(
                                            "BridgeLine successfully replaced: {:?}",
                                            bridgeline
                                        );
                                    } else if res == lox_library::ReplaceSuccess::NotReplaced {
                                        // Add the bridge to the list of to_be_replaced bridges in the Lox context and try
                                        // again to replace at the next update (nothing changes in the Lox Authority)
                                        println!(
                                    "'Gone' BridgeLine NOT replaced, saved for next update! : {:?}",
                                    bridgeline
                                );
                                        context.new_to_be_replaced_bridge(bridgeline);
                                    }
                                }
                            }
                        }
                    }
                }
                /* Functionality for marking bridges as unreachable/blocked is currently not enabled as there is not
                yet a reliable way to determine that a bridge is blocked. This means that migrations to unblocked bridges do not
                currently work but can be easily enabled with a list of `blocked resources` from rdsys or another source with something
                like the following:
                            println!("BridgeLine to be removed: {:?}", bridgeline);
                            let res = context.add_unreachable(bridgeline);
                            if res {
                                println!(
                                    "BridgeLine successfully marked unreachable: {:?}",
                                    bridgeline
                                );
                            } else {
                                println!("'Gone' BridgeLine NOT REMOVED!! : {:?}", bridgeline);
                                //TODO probably do something else here
                            }
                */
                context.allocate_leftover_bridges();
                context.encrypt_table();
                let mut date = Local::now().format("%Y-%m-%d_%H:%M:%S").to_string();
                let path = "_lox.json";
                date.push_str(path);
                let file = File::create(&date).expect("Unable to write to file!");
                let _ = serde_json::to_writer(file, &context);

                sleep(Duration::from_millis(1)).await;
            }
            Request { req, sender } => {
                let response = handle(context.clone(), req).await;
                if let Err(e) = sender.send(response) {
                    eprintln!("Server Response Error: {:?}", e);
                };
                sleep(Duration::from_millis(1)).await;
            }
            Shutdown { shutdown_sig } => {
                println!("Sending Shutdown Signal, all threads should shutdown.");
                drop(shutdown_sig);
                println!("Shutdown Sent.");
            }
        }
    }
}

// Each of the commands that the Context Manager handles
#[derive(Debug)]
enum Command {
    Rdsys {
        resourcediff: ResourceDiff,
    },
    Request {
        req: Request<Body>,
        sender: oneshot::Sender<Result<Response<Body>, Infallible>>,
    },
    Shutdown {
        shutdown_sig: broadcast::Sender<()>,
    },
}

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

    let file = File::open(&args.config).expect("Should have been able to read config.json file");
    let reader = BufReader::new(file);
    // Read the JSON contents of the file as a ResourceInfo
    let rtype: ResourceInfo =
        serde_json::from_reader(reader).expect("Reading ResourceInfo from JSON failed.");

    let (rdsys_tx, context_rx) = mpsc::channel(32);
    let request_tx = rdsys_tx.clone();
    let shutdown_cmd_tx = rdsys_tx.clone();

    // create the shutdown broadcast channel and clone for every thread
    let (shutdown_tx, mut shutdown_rx) = broadcast::channel(16);
    let kill_stream = shutdown_tx.subscribe();
    let kill_parser = shutdown_tx.subscribe();
    let kill_context = shutdown_tx.subscribe();

    // Listen for ctrl_c, send signal to broadcast shutdown to all threads by dropping shutdown_tx
    let shutdown_handler = spawn(async move {
        tokio::select! {
            _ = signal::ctrl_c() => {
                let cmd = Command::Shutdown {
                    shutdown_sig: shutdown_tx,
                };
                shutdown_cmd_tx.send(cmd).await.unwrap();
                sleep(Duration::from_secs(1)).await;

               _ = shutdown_rx.recv().await;
            }
        }
    });

    let context_manager =
        spawn(async move { create_context_manager(args.db, context_rx, kill_context).await });

    let (tx, rx) = mpsc::channel(32);
    let rdsys_stream_handler = spawn(async { rdsys_stream(rtype, tx, kill_stream).await });

    let rdsys_resource_receiver =
        spawn(async { rdsys_bridge_parser(rdsys_tx, rx, kill_parser).await });

    let make_service = make_service_fn(move |_conn: &AddrStream| {
        let request_tx = request_tx.clone();
        let service = service_fn(move |req| {
            let request_tx = request_tx.clone();
            let (response_tx, response_rx) = oneshot::channel();
            let cmd = Command::Request {
                req,
                sender: response_tx,
            };
            async move {
                request_tx.send(cmd).await.unwrap();
                response_rx.await.unwrap()
            }
        });
        async move { Ok::<_, Infallible>(service) }
    });

    let addr = SocketAddr::from(([127, 0, 0, 1], 8001));
    let server = Server::bind(&addr).serve(make_service);
    let graceful = server.with_graceful_shutdown(shutdown_signal());
    println!("Listening on {}", addr);
    if let Err(e) = graceful.await {
        eprintln!("server error: {}", e);
    }
    future::join_all([
        rdsys_stream_handler,
        rdsys_resource_receiver,
        context_manager,
        shutdown_handler,
    ])
    .await;
}