use std::{
	collections::{BTreeMap, btree_map},
	io::{Read, Write},
	sync::mpsc::{Receiver, Sender, channel},
};

const CLIENT_TO_SERVER: u8 = b'C';
const SERVER_TO_CLIENT: u8 = b'S';

pub type Records = BTreeMap<u64, (Vec<u8>, Vec<(Direction, Vec<u8>)>)>;

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Direction {
	ClientToServer,
	ServerToClient,
}

pub struct Recorder {
	file: std::fs::File,
	receiver: Receiver<(u64, Option<String>, Direction, Vec<u8>)>,
}

impl Recorder {
	#[allow(clippy::type_complexity)]
	pub fn new(path: &str) -> (Self, Sender<(u64, Option<String>, Direction, Vec<u8>)>) {
		let (sender, receiver) = channel();
		(
			Self {
				file: std::fs::OpenOptions::new()
					.create(true)
					.write(true)
					.truncate(true)
					.open(path)
					.unwrap(),
				receiver,
			},
			sender,
		)
	}

	pub fn run(&mut self) {
		while let Ok((conn_id, server_name, direction, data)) = self.receiver.recv() {
			let Some(server_name) = server_name else {
				continue;
			};
			let server_name = server_name.as_bytes();
			self.file
				.write_all(&[match direction {
					Direction::ClientToServer => CLIENT_TO_SERVER,
					Direction::ServerToClient => SERVER_TO_CLIENT,
				}])
				.unwrap();
			self.file.write_all(&conn_id.to_be_bytes()).unwrap();
			self.file.write_all(&[server_name.len() as u8]).unwrap();
			self.file.write_all(server_name).unwrap();
			self.file
				.write_all(&(data.len() as u64).to_be_bytes())
				.unwrap();
			self.file.write_all(&data).unwrap();
			self.file.flush().unwrap();
		}
	}
}

#[derive(Clone)]
struct Handler {
	sender: Sender<(u64, Option<String>, Direction, Vec<u8>)>,
	server_name: Option<String>,
}

impl sslrelay::HandlerCallbacks for Handler {
	// DownStream non blocking callback
	fn ds_nb_callback(&self, in_data: Vec<u8>, conn_id: u64) {
		self.sender
			.send((
				conn_id,
				self.server_name.clone(),
				Direction::ClientToServer,
				in_data,
			))
			.unwrap();
	}

	// DownStream blocking callback
	fn ds_b_callback(&mut self, in_data: Vec<u8>, _conn_id: u64) -> sslrelay::CallbackRet {
		sslrelay::CallbackRet::Relay(in_data)
	}

	// UpStream non blocking callback
	fn us_nb_callback(&self, in_data: Vec<u8>, conn_id: u64) {
		self.sender
			.send((
				conn_id,
				self.server_name.clone(),
				Direction::ServerToClient,
				in_data,
			))
			.unwrap();
	}

	// UpStream blocking callback
	fn us_b_callback(&mut self, in_data: Vec<u8>, _conn_id: u64) -> sslrelay::CallbackRet {
		sslrelay::CallbackRet::Relay(in_data)
	}

	fn set_server_name(&mut self, server_name: Option<&str>) {
		self.server_name = server_name.map(str::to_string);
	}
}

pub fn make_record(path: &str) {
	let (mut recorder, sender) = Recorder::new(path);
	let mut relay = sslrelay::SSLRelay::new(
		Handler {
			sender,
			server_name: None,
		},
		sslrelay::RelayConfig {
			downstream_data_type: sslrelay::TCPDataType::TLS,
			upstream_data_type: sslrelay::TCPDataType::TLS,
			bind_host: "127.0.0.1".to_string(),
			bind_port: "443".to_string(),
			remote_host: |server_name| {
				server_name
					.map(str::to_string)
					.unwrap_or_else(|| String::from("www.apple.com"))
			},
			remote_port: "443".to_string(),
			tls_config: sslrelay::TLSConfig::FILE {
				certificate_path: "/dev/shm/exp/certs/prime256v1/all.crt".to_string(),
				private_key_path: "/dev/shm/exp/certs/prime256v1/all.key".to_string(),
			},
		},
	);

	std::thread::spawn(move || recorder.run());

	relay.start();
}

pub fn read_record_file(path: &str) -> Records {
	let mut file = std::fs::OpenOptions::new().read(true).open(path).unwrap();
	let mut records = BTreeMap::<u64, (Vec<u8>, Vec<(Direction, Vec<u8>)>)>::new();
	loop {
		let mut direction = [0; 1];
		if file.read(&mut direction).unwrap() != 1 {
			break;
		}
		let direction = match direction[0] {
			CLIENT_TO_SERVER => Direction::ClientToServer,
			SERVER_TO_CLIENT => Direction::ServerToClient,
			_ => {
				println!("Error: invalid direction. stop.");
				break;
			}
		};
		let mut conn_id = [0; 8];
		if file.read(&mut conn_id).unwrap() != 8 {
			println!("Error: incomplete conn id. stop.");
			break;
		}
		let conn_id = u64::from_be_bytes(conn_id);
		let mut server_name_len = [0];
		if file.read(&mut server_name_len).unwrap() != 1 {
			println!("Error: incomplete server name len. stop.");
			break;
		}
		let server_name_len = server_name_len[0] as usize;
		let mut server_name = vec![0; server_name_len];
		if file.read(&mut server_name).unwrap() != server_name_len {
			println!("Error: incomplete data. stop.");
			break;
		}
		let mut len = [0; 8];
		if file.read(&mut len).unwrap() != 8 {
			println!("Error: incomplete len. stop.");
			break;
		}
		let len = u64::from_be_bytes(len);
		if len > 0xfff_ffff {
			println!("Error: len too large {len}. stop.");
			break;
		}
		let mut buf = vec![0; len as usize];
		if file.read(&mut buf).unwrap() != len as usize {
			println!("Error: incomplete data. stop.");
			break;
		}
		match records.entry(conn_id) {
			btree_map::Entry::Occupied(mut entry) => {
				entry.get_mut().1.push((direction, buf));
			}
			btree_map::Entry::Vacant(entry) => {
				entry.insert((server_name, vec![(direction, buf)]));
			}
		}
	}
	records
}

pub fn print_records(records: &Records, print_packets: bool) {
	for (id, (server_name, records)) in records {
		let server_name = str::from_utf8(server_name.as_slice()).unwrap();
		println!("{id} {server_name}");
		for (direction, data) in records {
			match direction {
				Direction::ClientToServer => {
					println!("  >>  {}", data.len());
				}
				Direction::ServerToClient => {
					println!("   << {}", data.len());
				}
			}
			if print_packets {
				let data = if data.len() >= 256 {
					&data[0..256]
				} else {
					data.as_slice()
				};
				if let Ok(data) = str::from_utf8(data) {
					println!("       {data:?}")
				} else {
					println!("       {data:?}")
				}
			}
		}
	}
}