use std::{cell::UnsafeCell, error::Error, net::Ipv4Addr, rc::Rc, thread::JoinHandle};

use count::Count;
use ethertype::EtherType;
use ippacket::{IpNextHeader, Ipv4Packet};
use layer3::{Tcp, Udp};
use pnet_datalink::{Channel, Config};
use scurvy::{Argument, Scurvy};
use smol::{
	LocalExecutor,
	channel::{self, Receiver, Sender},
	io::AsyncWriteExt,
	net::TcpListener,
};

mod count;
mod ethertype;
mod ippacket;
mod layer3;

fn main() -> Result<(), Box<dyn Error>> {
	let args = vec![
		Argument::new(&["interface", "iface"]).arg("dev"),
		Argument::new("ip").arg("addr"),
	];
	let scurvy = Scurvy::make(args);

	let interface_name = scurvy.get("interface").unwrap();
	let interface = match pnet_datalink::interfaces().iter().find(|i| i.name == interface_name) {
		None => {
			eprintln!("No interface found named '{interface_name}'");
			return Ok(());
		}
		Some(i) => i.clone(),
	};

	let ip_want: Ipv4Addr = scurvy.get("ip").unwrap().parse()?;

	let mut channel = match pnet_datalink::channel(&interface, Config::default())? {
		Channel::Ethernet(_tx, rx) => rx,
		_ => unimplemented!(),
	};

	let stat = stat_thread();

	loop {
		let pkt = channel.next()?;

		let ethertype = EtherType::new(u16::from_be_bytes([pkt[12], pkt[13]]));
		let eth_payload = &pkt[14..];

		match ethertype {
			EtherType::IPv4 => {
				let ip = Ipv4Packet::new(eth_payload);
				let ip_payload = ip.get_payload();

				let ip_tx = if ip.src == ip_want {
					true
				} else if ip.dst == ip_want {
					false
				} else {
					continue;
				};

				// 6 byte per MAC (x2), 2 byte ethertype, 2 byte crc
				let total_l2_len = ip.get_packet_len() + 18;

				match ip.get_next_header() {
					IpNextHeader::Tcp => {
						let tcp = Tcp::new(ip_payload);

						stat.tx.send_blocking(Meow::Tcp {
							tx: ip_tx,
							src: tcp.source_port(),
							dst: tcp.destination_port(),
							len: total_l2_len,
						})?;
					}
					IpNextHeader::Udp => {
						let udp = Udp::new(ip_payload);

						stat.tx.send_blocking(Meow::Udp {
							tx: ip_tx,
							src: udp.source_port(),
							dst: udp.destination_port(),
							len: total_l2_len,
						})?;
					}
					_ => (),
				}
			}
			_ => (),
		}
	}
}

#[allow(dead_code)]
struct StatHandle {
	hwnd: JoinHandle<()>,
	tx: Sender<Meow>,
}

enum Meow {
	Tcp {
		tx: bool,
		src: u16,
		dst: u16,
		len: usize,
	},
	Udp {
		tx: bool,
		src: u16,
		dst: u16,
		len: usize,
	},
}

fn stat_thread() -> StatHandle {
	let (tx, rx) = channel::unbounded();

	let hwnd = std::thread::spawn(|| stat(rx));

	StatHandle { hwnd, tx }
}

fn stat(rx: Receiver<Meow>) {
	let cat = Rc::new(UnsafeCell::new(Count::new()));

	let ex = Rc::new(LocalExecutor::new());

	ex.spawn(async {
		loop {
			match rx.recv().await {
				Err(_e) => {
					eprintln!("error receiving! breaking from loop");
					break;
				}
				Ok(Meow::Tcp { tx, src, dst, len }) => {
					unsafe { &mut *cat.get() }.push_tcp(tx, src, dst, len);
				}
				Ok(Meow::Udp { tx, src, dst, len }) => {
					unsafe { &mut *cat.get() }.push_udp(tx, src, dst, len);
				}
			}
		}
	})
	.detach();

	let run_cat = Rc::clone(&cat);
	let run_ex = Rc::clone(&ex);
	let run = ex.run(async move {
		let socket = TcpListener::bind("127.0.0.1:6666").await.unwrap();

		loop {
			let (mut stream, _caddr) = socket.accept().await.unwrap();
			let per_spawn_cat = Rc::clone(&run_cat);

			run_ex
				.spawn(async move {
					let cat = unsafe { &mut *per_spawn_cat.get() };
					let http_tx = cat.many_tcp_tx(&[80, 443]);
					let http_rx = cat.many_tcp_rx(&[80, 443]);

					let str = format!("{http_tx}\n{http_rx}");
					stream.write_all(str.as_bytes()).await.unwrap();
				})
				.detach();
		}
	});

	smol::block_on(run);
}