// SPDX-License-Identifier: EUPL-1.2+
// SPDX-FileCopyrightText: 2025 Yureka Lilian <yureka@cyberchaos.dev>

use std::collections::HashMap;
use std::io;
use std::net::Ipv6Addr;
use std::path::PathBuf;

use clap::Parser;
use futures_util::SinkExt;
use futures_util::StreamExt;
use log::{debug, error, info, trace, warn};
use tokio::net::UnixListener;
use tokio_stream::StreamMap;
use vhost_device_net::VhostDeviceNet;

use zerocopy::byteorder::network_endian::{U16, U32};
use zerocopy::*;

type MacAddr = [u8; 6];

#[derive(Debug, PartialEq, Eq, FromBytes, IntoBytes, KnownLayout, Immutable, Unaligned)]
#[repr(C)]
struct EtherFrame {
    //preamble_sfd: U64,
    dst_addr: MacAddr,
    src_addr: MacAddr,
    ether_type: U16,
    //data: [u8]
}

impl AsRef<[u8]> for EtherFrame {
    fn as_ref(&self) -> &[u8] {
        self.as_bytes()
    }
}

#[derive(Debug, PartialEq, Eq, FromBytes, IntoBytes, KnownLayout, Immutable, Unaligned)]
#[repr(C)]
struct VlanTag {
    tag_control_information: U16,
    ether_type: U16,
}

impl AsRef<[u8]> for VlanTag {
    fn as_ref(&self) -> &[u8] {
        self.as_bytes()
    }
}

#[derive(Debug, PartialEq, Eq, FromBytes, IntoBytes, KnownLayout, Immutable, Unaligned)]
#[repr(C)]
struct Ipv6Header {
    version_traffic_class_flow_label: U32,
    payload_length: U16,
    hext_header: u8,
    hop_limit: u8,
    src_addr: [u8; 16],
    dst_addr: [u8; 16],
    //data: [u8]
}

impl AsRef<[u8]> for Ipv6Header {
    fn as_ref(&self) -> &[u8] {
        self.as_bytes()
    }
}

#[derive(Parser, Debug)]
#[command()] //version = None, about = None, long_about = None)]
struct Args {
    #[arg(long)]
    driver_listen_path: PathBuf,
    #[arg(long)]
    app_listen_path: PathBuf,
    #[arg(long)]
    upstream_interface: u16,
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
    env_logger::init();
    let args = Args::parse();

    run_router(args).await
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum InterfaceId {
    Driver,
    App(usize),
}

async fn run_router(args: Args) -> anyhow::Result<()> {
    let driver_listener = UnixListener::bind(&args.driver_listen_path)?;
    let app_listener = UnixListener::bind(&args.app_listen_path)?;

    let fib: HashMap<Ipv6Addr, (MacAddr, InterfaceId)> = Default::default();

    let mut fib = fib.clone();
    let mut phys_mac = None;

    let mut streams = StreamMap::new();
    let mut sinks = HashMap::<InterfaceId, _>::new();

    let mut app_num = 0;

    loop {
        tokio::select! {
            driver_conn = driver_listener.accept() => {
                info!("driver connected");
                match driver_conn {
                    Ok((stream, _addr)) => {
                        let device = VhostDeviceNet::from_unix_stream(stream).await?;
                        streams.insert(InterfaceId::Driver, Box::pin(device.tx().await?));
                        sinks.insert(InterfaceId::Driver, Box::pin(device.rx::<Box<dyn io::Read + Send + Sync + 'static>>().await?));
                        phys_mac = None;
                    }
                    Err(e) => error!("driver connection failed: {}", e),
                }
            }
            app_conn = app_listener.accept() => {
                info!("app connected");
                match app_conn {
                    Ok((stream, _addr)) => {
                        let device = VhostDeviceNet::from_unix_stream(stream).await?;
                        streams.insert(InterfaceId::App(app_num), Box::pin(device.tx().await?));
                        sinks.insert(InterfaceId::App(app_num), Box::pin(device.rx().await?));
                        app_num = app_num.checked_add(1).unwrap();
                    }
                    Err(e) => error!("app connection failed: {}", e),
                }
            }
            next_res = streams.next(), if !streams.is_empty() => {
                let Some((key, Ok(buf))) = next_res else {
                    info!("incoming other");
                    continue;
                };
                match &key {
                    InterfaceId::Driver => {
                        use io::Read;
                        let mut ether_bytes = buf.take(size_of::<EtherFrame>() as u64);
                        let mut ether_frame = EtherFrame::read_from_io(&mut ether_bytes)?;
                        let remaining_buf = ether_bytes.into_inner();

                        if ether_frame.ether_type != 0x8100 {
                            warn!("untagged packet from driver");
                            //eprintln!("{:x?}", ether_frame);
                            //let mut ipv6_bytes = remaining_buf.take(size_of::<Ipv6Header>() as u64);
                            //let ipv6_hdr = Ipv6Header::read_from_io(&mut ipv6_bytes)?;
                            //let remaining_buf = ipv6_bytes.into_inner();
                            //eprintln!("{:x?}", ipv6_hdr);

                            continue;
                        }

                        let mut vlan_bytes = remaining_buf.take(size_of::<VlanTag>() as u64);
                        let vlan_tag = VlanTag::read_from_io(&mut vlan_bytes)?;
                        let remaining_buf = vlan_bytes.into_inner();

                        let vlan_id = u16::from(vlan_tag.tag_control_information) & 0xfff;
                        if vlan_id != args.upstream_interface {
                            debug!("dropping packet with vlan {}", vlan_id);
                            continue;
                        }
                        if !phys_mac.is_some() {
                            debug!("set phys mac");
                            phys_mac = Some(ether_frame.dst_addr);
                        }

                        if vlan_tag.ether_type != 0x86dd {
                            continue;
                        }
                        let mut ipv6_bytes = remaining_buf.take(size_of::<Ipv6Header>() as u64);
                        let ipv6_hdr = Ipv6Header::read_from_io(&mut ipv6_bytes)?;
                        let remaining_buf = ipv6_bytes.into_inner();
                        trace!("{:x?}", ipv6_hdr);
                        let dst_addr = Ipv6Addr::from(ipv6_hdr.dst_addr);

                        let Some((dst_mac, if_idx)) = fib.get(&dst_addr) else {
                            warn!("dropped incoming message to {} because no fib match", dst_addr);
                            continue;
                        };
                        ether_frame.dst_addr = *dst_mac;
                        ether_frame.ether_type = vlan_tag.ether_type;

                        sinks.get_mut(if_idx).unwrap().send(Box::new(io::Cursor::new(ether_frame).chain(io::Cursor::new(ipv6_hdr)).chain(remaining_buf))).await.unwrap();
                    }
                    InterfaceId::App(_) => {
                        use io::Read;
                        let mut ether_bytes = buf.take(size_of::<EtherFrame>() as u64);
                        let mut ether_frame = EtherFrame::read_from_io(&mut ether_bytes)?;
                        let remaining_buf = ether_bytes.into_inner();
                        trace!("{:x?}", ether_frame);
                        if ether_frame.ether_type != 0x86dd {
                            continue;
                        }
                        let mut ipv6_bytes = remaining_buf.take(size_of::<Ipv6Header>() as u64);
                        let ipv6_hdr = Ipv6Header::read_from_io(&mut ipv6_bytes)?;
                        let remaining_buf = ipv6_bytes.into_inner();
                        trace!("{:x?}", ipv6_hdr);
                        let src_addr = Ipv6Addr::from(ipv6_hdr.src_addr);
                        trace!("{:x?}", src_addr);
                        if !src_addr.is_unspecified() && !src_addr.is_multicast() && !fib.contains_key(&src_addr) {
                            debug!("adding fib entry for {}", src_addr);
                            fib.insert(src_addr, (ether_frame.src_addr.clone(), key));
                        }

                        let Some(phys_mac) = &phys_mac else {
                            warn!("dropped message because phys mac is unknown");
                            continue;
                        };
                        let vlan_tag = VlanTag {
                            ether_type: ether_frame.ether_type,
                            tag_control_information: args.upstream_interface.into(),
                        };
                        ether_frame.src_addr = phys_mac.clone();
                        ether_frame.ether_type = 0x8100.into();

                        let Some(sink) = sinks.get_mut(&InterfaceId::Driver) else {
                            warn!("dropped message because driver is not ready");
                            continue;
                        };

                        sink.send(Box::new(io::Cursor::new(ether_frame)
                            .chain(io::Cursor::new(vlan_tag))
                            .chain(io::Cursor::new(ipv6_hdr))
                            .chain(remaining_buf)))
                            .await?;
                    }
                }
            }
        }
    }
}