#![feature(let_chains)]

mod client;
mod http;
mod record;
mod server;

use record::Records;

use argp::FromArgs;
use static_cell::StaticCell;
use tokio::sync::oneshot;

/// Play recorded requests and responses
#[derive(FromArgs)]
struct Opt {
	/// Path to record file
	#[argp(positional)]
	record_file: String,
	#[argp(subcommand)]
	subcommand: Subcommand,
}

#[derive(FromArgs)]
#[argp(subcommand)]
enum Subcommand {
	/// Replay from records
	Play(OptPlay),
	/// Print records
	Print(OptPrint),
	/// Record traffic
	Record(OptRecord),
}

/// Replay from records
#[derive(FromArgs)]
#[argp(subcommand, name = "play")]
struct OptPlay {
	/// Connect to address
	#[argp(positional)]
	forward_addr: String,
	/// Connect to port
	#[argp(positional)]
	forward_port: u16,
	/// Listen to port
	#[argp(positional)]
	listen_port: u16,
	/// Path to PEM certificates and keys
	#[argp(positional)]
	certs: String,
	/// Where to use TLS
	#[argp(positional)]
	tls: String,
	/// Repeat N times
	#[argp(option, short = 'r', default = "1")]
	repeat: u32,
	/// Only play this record
	#[argp(option)]
	record: Option<u64>,
}

/// Print records
#[derive(FromArgs)]
#[argp(subcommand, name = "print")]
struct OptPrint {
	/// Print packets
	#[argp(switch, short = 'p')]
	packets: bool,
}

/// Record traffic
#[derive(FromArgs)]
#[argp(subcommand, name = "record")]
struct OptRecord {}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TlsMode {
	None,
	Client,
	Server,
	Both,
}

static RECORDS: StaticCell<Records> = StaticCell::new();

#[tokio::main]
async fn main() {
	let opt: Opt = argp::parse_args_or_exit(argp::DEFAULT);

	match opt.subcommand {
		Subcommand::Play(subopt) => {
			let tls_mode = match subopt.tls.as_str() {
				"none" => TlsMode::None,
				"client" => TlsMode::Client,
				"server" => TlsMode::Server,
				"both" => TlsMode::Both,
				_ => panic!("TLS mode must be one of none,client,server,both."),
			};
			let records = RECORDS.init(record::read_record_file(&opt.record_file));

			if let Some(only_record) = subopt.record {
				records.retain(|id, _| *id == only_record);
			}

			let (sync_sender, sync_receiver) = oneshot::channel();
			console_subscriber::init();
			let client = tokio::spawn(client::play(
				records,
				tls_mode,
				(subopt.forward_addr, subopt.forward_port),
				sync_receiver,
				subopt.repeat,
			));
			server::play(
				records,
				tls_mode,
				&subopt.certs,
				("0.0.0.0", subopt.listen_port),
				sync_sender,
			)
			.await;
			client.await.unwrap();
		}
		Subcommand::Print(subopt) => {
			let records = record::read_record_file(&opt.record_file);
			record::print_records(&records, subopt.packets);
		}
		Subcommand::Record(_subopt) => {
			record::make_record(&opt.record_file);
		}
	}
}