mod datastore;
mod grpc;
mod http_server;
mod persistence;
mod solana;

use crate::{
    datastore::LOCALHOST, grpc::challenge::NodeUpdate, persistence::KeysFile, solana::SolClient,
};
use datastore::State;
use detee_sgx::{InstanceMeasurement, RaTlsConfig};
use std::{
    fs::File,
    io::{BufRead, BufReader},
    sync::Arc,
};
use tokio::{
    sync::{broadcast, broadcast::Sender},
    task::JoinSet,
    time::{sleep, Duration},
};

const INIT_NODES_FILE: &str = "/host/detee_challenge_nodes";
const KEYS_FILE: &str = "/host/main/TRY_TO_HACK_THIS";
const MAX_CONNECTIONS: usize = 3;

pub async fn localhost_cron(state: Arc<State>, tx: Sender<NodeUpdate>) {
    loop {
        sleep(Duration::from_secs(60)).await;
        let _ = tx.send((LOCALHOST.to_string(), state.get_localhost()).into());
        state.remove_inactive_nodes();
    }
}

async fn get_sol_client(state: Arc<State>, ratls_config: RaTlsConfig) -> SolClient {
    match KeysFile::read(KEYS_FILE, &state).await {
        Ok(keys_file) => {
            let sol_client = SolClient::try_from(keys_file).unwrap();
            println!(
                "Found the following wallet saved to disk: {}",
                sol_client.get_wallet_pubkey()
            );
            println!("Loading token mint address {}", sol_client.get_token_address());
            return sol_client;
        }
        Err(e) => println!("Can't initialize using sealed keys: {e}"),
    };

    let init_nodes = match File::open(INIT_NODES_FILE) {
        Ok(init_nodes) => init_nodes,
        Err(_) => {
            println!("Can't initialize using init nodes from {INIT_NODES_FILE}");
            println!("Starting a new network with a new key...");
            return SolClient::new().await;
        }
    };

    let init_nodes_reader = BufReader::new(init_nodes);
    for init_node_ip in init_nodes_reader.lines().map(|l| l.unwrap()) {
        match grpc::client::key_grabber(state.clone(), init_node_ip, ratls_config.clone()).await {
            Ok(keys) => {
                let sol_client = SolClient::try_from(keys.clone())
                    .map_err(|e| {
                        println!("Received malformed keys from the network: {e}");
                        state.increase_net_attacks();
                    })
                    .unwrap();
                println!(
                    "Got keypair from the network. Joining the network using wallet {}",
                    sol_client.get_wallet_pubkey()
                );
                println!("The address of the Token is {}", sol_client.get_token_address());
                println!("Saving this data to disk in the file {KEYS_FILE}");
                if let Err(e) = sol_client.get_keys_file().write(KEYS_FILE).await {
                    println!("Could not save data to disk: {e}");
                }
                return sol_client;
            }
            Err(e) => {
                println!("Could not get keypair: {e:?}");
            }
        };
    }
    panic!("could not get keypair.");
}

#[tokio::main]
async fn main() {
    env_logger::init_from_env(env_logger::Env::default().default_filter_or("warn"));
    let ratls_config = RaTlsConfig::new()
        .allow_instance_measurement(InstanceMeasurement::new().with_current_mrenclave().unwrap());

    let state = Arc::new(State::new());
    let sol_client = Arc::new(get_sol_client(state.clone(), ratls_config.clone()).await);

    let (tx, _) = broadcast::channel(500);

    let mut tasks = JoinSet::new();

    tasks.spawn(localhost_cron(state.clone(), tx.clone()));
    tasks.spawn(http_server::init(state.clone(), sol_client.clone()));
    tasks.spawn(
        grpc::server::MyServer::init(
            state.clone(),
            sol_client.get_keys(),
            ratls_config.clone(),
            tx.clone(),
        )
        .start(),
    );

    if let Ok(input) = std::fs::read_to_string(INIT_NODES_FILE) {
        for line in input.lines() {
            tasks.spawn(
                grpc::client::ConnManager::init(state.clone(), ratls_config.clone(), tx.clone())
                    .start_with_node(line.to_string()),
            );
        }
    }

    for _ in 0..MAX_CONNECTIONS {
        tasks.spawn(
            grpc::client::ConnManager::init(state.clone(), ratls_config.clone(), tx.clone())
                .start(),
        );
    }

    while let Some(Ok(_)) = tasks.join_next().await {}

    // task panicked
    println!("Shutting down...");
}