use core::{
    fmt,
    net::IpAddr,
    task::{Context, Poll},
    time::Duration,
};

use alloc::{boxed::Box, sync::Arc, vec::Vec};
use async_trait::async_trait;
use libk::{
    error::Error,
    vfs::{FileReadiness, Socket},
};
use libk_mm::PageBox;
use libk_util::queue::BoundedMpmcQueue;
use yggdrasil_abi::{
    abi_serde::{wire, Serialize},
    net::{
        netconfig::{
            InterfaceBinding, InterfaceInfo, NetConfigRequest, NetConfigResult, RouteInfo,
            RoutingInfo,
        },
        MacAddress, MessageHeader, MessageHeaderMut, SocketInterfaceQuery, SubnetAddr,
    },
};

use crate::{
    interface::NetworkInterface,
    l3::{arp, Route},
};

pub struct NetConfigSocket {
    // Responses are written here
    buffer: BoundedMpmcQueue<(usize, PageBox<[u8; 4096]>)>,
}

impl NetConfigSocket {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            buffer: BoundedMpmcQueue::new(8),
        })
    }

    fn write_response<T: Serialize>(&self, value: NetConfigResult<T>) -> Result<(), Error> {
        let mut response = PageBox::new([0; 4096])?;
        let len = wire::to_slice(&value, &mut response[..]).unwrap();
        self.buffer
            .try_push_back((len, response))
            .map_err(|_| Error::OutOfMemory)
    }

    fn read_response(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        // Don't allow truncated responses
        let (len, data) = self.buffer.try_pop_front().ok_or(Error::WouldBlock)?;
        if len > buffer.len() {
            log::warn!(
                "Buffer of {} is too small to write a response of {len}",
                buffer.len()
            );
            return Err(Error::InvalidArgument);
        }
        buffer[..len].copy_from_slice(&data[..len]);
        Ok(len)
    }

    async fn handle_request(&self, request: NetConfigRequest<'_>) -> Result<(), Error> {
        match request {
            NetConfigRequest::ArpQuery(query) if !query.cache_only => {
                let iface = match query_interface(query.interface) {
                    Some(iface) => iface,
                    None => {
                        return self.write_response(NetConfigResult::<MacAddress>::Err(
                            "No such interface",
                        ))
                    }
                };
                let result = match arp::lookup(iface.id(), query.address, true).await {
                    Ok(mac) => NetConfigResult::Ok(mac),
                    Err(_) => NetConfigResult::Err("No ARP entry"),
                };
                self.write_response(result)
            }
            _ => self.handle_request_nonblocking(request),
        }
    }

    fn handle_request_nonblocking(&self, request: NetConfigRequest) -> Result<(), Error> {
        match request {
            NetConfigRequest::ListRoutes(()) => {
                let routes = Route::list_ref();
                let route_info: Vec<_> = routes.iter().map(describe_route).collect();
                self.write_response(NetConfigResult::Ok(route_info))
            }
            NetConfigRequest::ListInterfaces(()) => {
                let interfaces = NetworkInterface::list_ref();
                let interfaces: Vec<_> = interfaces
                    .iter()
                    .map(|(&id, iface)| InterfaceBinding {
                        name: iface.name.clone(),
                        id,
                    })
                    .collect();
                self.write_response(NetConfigResult::Ok(interfaces))
            }
            NetConfigRequest::DescribeRoutes(_query) => todo!(),
            NetConfigRequest::DescribeInterface(query) => {
                let result = match query_interface(query) {
                    Some(interface) => NetConfigResult::Ok(describe_interface(&interface)),
                    None => NetConfigResult::Err("No such interface"),
                };
                self.write_response(result)
            }
            NetConfigRequest::AddRoute(op) => {
                let result = match add_route(op.interface, op.gateway, op.subnet) {
                    Ok(()) => NetConfigResult::Ok(()),
                    Err(error) => NetConfigResult::Err(error),
                };
                self.write_response(result)
            }
            NetConfigRequest::ArpQuery(query) => {
                let iface = match query_interface(query.interface) {
                    Some(iface) => iface,
                    None => {
                        return self.write_response(NetConfigResult::<MacAddress>::Err(
                            "No such interface",
                        ))
                    }
                };
                let result =
                    match arp::lookup_nonblocking(iface.id(), query.address, !query.cache_only) {
                        Ok(mac) => NetConfigResult::Ok(mac),
                        Err(_) => NetConfigResult::Err("No ARP entry"),
                    };
                self.write_response(result)
            }
            NetConfigRequest::QueryRoute(query) => {
                let result = match query_route(query) {
                    Some(route) => NetConfigResult::Ok(route),
                    None => NetConfigResult::Err("No route to host"),
                };
                self.write_response(result)
            }
            NetConfigRequest::SetNetworkAddress(op) => {
                let result = match query_interface(op.interface) {
                    Some(interface) => {
                        interface.set_address(op.address);
                        NetConfigResult::Ok(())
                    }
                    None => NetConfigResult::Err("No such interface"),
                };
                self.write_response(result)
            }
            NetConfigRequest::ClearNetworkAddress(_) => todo!(),
        }
    }
}

#[async_trait]
impl Socket for NetConfigSocket {
    fn bind(self: Arc<Self>, _local: &[u8]) -> Result<(), Error> {
        Err(Error::InvalidOperation)
    }
    async fn connect(
        self: Arc<Self>,
        _remote: &[u8],
        _timeout: Option<Duration>,
    ) -> Result<(), Error> {
        Err(Error::InvalidOperation)
    }

    fn close(self: Arc<Self>) -> Result<(), Error> {
        Ok(())
    }

    fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        let _ = option;
        let _ = buffer;
        Err(Error::InvalidOperation)
    }

    fn get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        let _ = option;
        let _ = buffer;
        Err(Error::InvalidOperation)
    }

    fn send_nonblocking(self: Arc<Self>, message: &MessageHeader) -> Result<usize, Error> {
        let request: NetConfigRequest = wire::from_slice(message.payload)
            .inspect_err(|_| log::warn!("Invalid netconfig request"))?;
        self.handle_request_nonblocking(request)?;
        Ok(message.payload.len())
    }
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        _timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let request: NetConfigRequest = wire::from_slice(message.payload)
            .inspect_err(|_| log::warn!("Invalid netconfig request"))?;
        self.handle_request(request).await?;
        Ok(message.payload.len())
    }

    fn receive_nonblocking(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
    ) -> Result<usize, Error> {
        self.read_response(message.payload)
    }
    async fn receive_message(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
        _timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        self.read_response(message.payload)
    }
}

impl FileReadiness for NetConfigSocket {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.buffer.poll_not_empty(cx).map(|_| Ok(()))
    }
}

impl fmt::Debug for NetConfigSocket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("NetConfigSocket").finish_non_exhaustive()
    }
}

fn query_interface(query: SocketInterfaceQuery) -> Option<Arc<NetworkInterface>> {
    match query {
        SocketInterfaceQuery::ById(id) => NetworkInterface::get(id).ok(),
        SocketInterfaceQuery::ByName(name) => {
            let interfaces = NetworkInterface::list_ref();
            interfaces.iter().find_map(|(_, iface)| {
                if &*iface.name == name {
                    Some(iface.clone())
                } else {
                    None
                }
            })
        }
    }
}

fn describe_interface(interface: &NetworkInterface) -> InterfaceInfo {
    InterfaceInfo {
        interface_id: interface.id,
        interface_name: interface.name.clone(),
        address: interface.address.read().map(Into::into),
        mac: interface.mac,
        link: interface.link_state(),
    }
}

fn describe_route(route: &Route) -> RouteInfo {
    // NOTE: must exist
    let interface = NetworkInterface::get(route.interface).unwrap();

    RouteInfo {
        interface_name: interface.name.clone(),
        interface_id: route.interface,
        subnet: route.subnet,
        gateway: route.gateway,
    }
}

fn add_route(
    query: SocketInterfaceQuery,
    gateway: Option<IpAddr>,
    subnet: SubnetAddr,
) -> Result<(), &'static str> {
    let interface = query_interface(query).ok_or("No such interface")?;
    let route = Route {
        interface: interface.id,
        gateway,
        subnet,
        priority: 10,
    };
    Route::insert(route).map_err(|_| "Could not insert route")?;
    Ok(())
}

fn query_route(destination: IpAddr) -> Option<RoutingInfo> {
    let (interface_id, gateway, destination) = Route::lookup(destination)?;
    let interface = NetworkInterface::get(interface_id).unwrap();
    let source = *interface.address.read();

    Some(RoutingInfo {
        interface_name: interface.name.clone(),
        interface_id,
        destination,
        gateway,
        source,
        source_mac: interface.mac,
    })
}
