use core::{
    fmt,
    sync::atomic::{AtomicU32, Ordering},
    task::{Context, Poll},
    time::Duration,
};

use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
use async_trait::async_trait;
use libk::{
    dma::DmaBuffer,
    error::Error,
    task::runtime::maybe_timeout,
    vfs::{FileReadiness, Socket},
};
use libk_mm::PageBox;
use libk_util::{queue::BoundedMpmcQueue, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::{
    net::{
        options::{self, RawSocketOptionVariant},
        MessageHeader, MessageHeaderMut, SocketInterfaceQuery,
    },
    option::OptionValue,
};

use crate::{ethernet::L2Packet, interface::NetworkInterface, TxPacketBuilder};

enum RawPacket {
    Ingress(L2Packet),
    Egress(Arc<PageBox<[u8]>>),
}

pub struct RawSocket {
    id: u32,
    bound: IrqSafeRwLock<Option<u32>>,
    receive_queue: BoundedMpmcQueue<RawPacket>,
}

static RAW_SOCKET_ID: AtomicU32 = AtomicU32::new(0);
static RAW_SOCKETS: IrqSafeRwLock<BTreeMap<u32, Arc<RawSocket>>> =
    IrqSafeRwLock::new(BTreeMap::new());
static BOUND_RAW_SOCKETS: IrqSafeRwLock<BTreeMap<u32, Vec<u32>>> =
    IrqSafeRwLock::new(BTreeMap::new());

impl RawSocket {
    pub fn new() -> Arc<Self> {
        let id = RAW_SOCKET_ID.fetch_add(1, Ordering::SeqCst);
        let socket = Arc::new(Self {
            id,
            bound: IrqSafeRwLock::new(None),
            receive_queue: BoundedMpmcQueue::new(256),
        });
        RAW_SOCKETS.write().insert(id, socket.clone());
        socket
    }

    fn bound_packet_received(&self, packet: RawPacket) {
        // TODO do something with the dropped packet?
        self.receive_queue.try_push_back(packet).ok();
    }

    pub fn handle_rx(packet: L2Packet) {
        let bound_sockets = BOUND_RAW_SOCKETS.read();
        let raw_sockets = RAW_SOCKETS.read();

        if let Some(ids) = bound_sockets.get(&packet.interface_id) {
            for id in ids {
                let socket = raw_sockets.get(id).unwrap();
                socket.bound_packet_received(RawPacket::Ingress(packet.clone()));
            }
        }
    }

    pub fn handle_tx(iface: u32, packet: &DmaBuffer<[u8]>, len: usize, except: Option<u32>) {
        let bound_sockets = BOUND_RAW_SOCKETS.read();
        let raw_sockets = RAW_SOCKETS.read();

        if let Some(ids) = bound_sockets.get(&iface) {
            let Ok(egress) = PageBox::from_slice(&packet[..len]).map(Arc::new) else {
                return;
            };

            for id in ids {
                if except == Some(*id) {
                    continue;
                }

                let socket = raw_sockets.get(id).unwrap();
                socket.bound_packet_received(RawPacket::Egress(egress.clone()));
            }
        }
    }

    fn packet_to_user(packet: RawPacket, buffer: &mut [u8]) -> Result<usize, Error> {
        match packet {
            RawPacket::Egress(egress) => {
                // TODO l2offset in egress?
                let len = egress.len();
                if buffer.len() < len {
                    return Err(Error::BufferTooSmall);
                }
                buffer[..len].copy_from_slice(&egress[..]);
                Ok(len)
            }
            RawPacket::Ingress(ingress) => {
                let full_len = ingress.data.len();
                let len = full_len - ingress.l2_offset;
                if buffer.len() < len {
                    return Err(Error::BufferTooSmall);
                }
                buffer[..len].copy_from_slice(&ingress.data[ingress.l2_offset..full_len]);
                Ok(len)
            }
        }
    }
}

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

#[async_trait]
impl Socket for RawSocket {
    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 get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        let option = RawSocketOptionVariant::try_from(option)?;
        match option {
            RawSocketOptionVariant::BoundInterface => {
                let bound = self.bound.read().ok_or(Error::NotConnected)?;
                options::BoundInterface::store(&SocketInterfaceQuery::ById(bound), buffer)
            }
            RawSocketOptionVariant::UnbindInterface => Err(Error::InvalidArgument),
            RawSocketOptionVariant::BoundHardwareAddress => {
                let bound = self.bound.read().ok_or(Error::NotConnected)?;
                let interface = NetworkInterface::get(bound)?;
                options::BoundHardwareAddress::store(&interface.mac, buffer)
            }
        }
    }

    fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        let option = RawSocketOptionVariant::try_from(option)?;
        match option {
            RawSocketOptionVariant::BoundInterface => {
                let query = options::BoundInterface::load(buffer)?;
                log::info!("raw: bind interface {query:?}");
                let mut bound = self.bound.write();
                if bound.is_some() {
                    return Err(Error::AlreadyExists);
                }

                let mut bound_sockets = BOUND_RAW_SOCKETS.write();
                let interface = match query {
                    SocketInterfaceQuery::ById(id) => NetworkInterface::get(id),
                    SocketInterfaceQuery::ByName(name) => NetworkInterface::query_by_name(name),
                }?;
                let list = bound_sockets.entry(interface.id).or_default();
                bound.replace(interface.id);
                list.push(self.id);

                Ok(())
            }
            RawSocketOptionVariant::UnbindInterface => {
                log::warn!("TODO: raw socket interface unbind");
                Err(Error::NotImplemented)
            }
            RawSocketOptionVariant::BoundHardwareAddress => Err(Error::ReadOnly),
        }
    }

    fn close(self: Arc<Self>) -> Result<(), Error> {
        let bound = self.bound.write().take();

        if let Some(bound) = bound {
            let mut bound_sockets = BOUND_RAW_SOCKETS.write();
            let mut clear = false;

            if let Some(list) = bound_sockets.get_mut(&bound) {
                list.retain(|&item| item != self.id);
                clear = list.is_empty();
            }

            if clear {
                bound_sockets.remove(&bound);
            }
        }

        RAW_SOCKETS.write().remove(&self.id).unwrap();

        Ok(())
    }

    fn send_nonblocking(self: Arc<Self>, message: &MessageHeader) -> Result<usize, Error> {
        // TODO cap by MTU?
        let bound = self.bound.read().ok_or(Error::InvalidOperation)?;
        let interface = NetworkInterface::get(bound)?;
        // let l2_offset = interface.device.packet_prefix_size();
        if message.payload.len() > 1024 {
            return Err(Error::InvalidArgument);
        }
        let mut builder = TxPacketBuilder::new(&interface, message.payload.len())?;
        builder.push_bytes(message.payload)?;
        // false to prevent loopback
        builder.transmit(Some(self.id))?;
        Ok(message.payload.len())
    }
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let _ = timeout;
        self.send_nonblocking(message)
    }

    fn receive_nonblocking(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
    ) -> Result<usize, Error> {
        let packet = self
            .receive_queue
            .try_pop_front()
            .ok_or(Error::WouldBlock)?;
        Self::packet_to_user(packet, message.payload)
    }
    async fn receive_message(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let future = self.receive_queue.pop_front();
        let packet = maybe_timeout(future, timeout).await?;
        Self::packet_to_user(packet, message.payload)
    }
}

impl fmt::Debug for RawSocket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let bound = *self.bound.read();
        f.debug_struct("RawSocket")
            .field("interface", &bound)
            .finish_non_exhaustive()
    }
}
