use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::time::Instant;

use dope::manifold::Manifold;
use dope::manifold::datagram::{Handler, Socket};
use dope::{Cqe, Drive, Driver};
use pin_project::pin_project;
use shin::sig::SigningKey;

use crate::conn::{Conn, ConnConfig, ConnHandle, DatagramError};
use crate::mux::{self, Mux};
use crate::transport_params;

#[pin_project]
pub struct Endpoint<const ID: u8, H: mux::Handler> {
    #[pin]
    udp: Socket<0>,
    mux: Mux<H>,
}

impl<const ID: u8, H: mux::Handler> Endpoint<ID, H> {
    pub fn build_server(
        bind: SocketAddr,
        signing_key: SigningKey,
        server_tp: transport_params::Params,
        handler: H,
        driver: &mut Driver,
    ) -> io::Result<Self> {
        Self::build_server_with_config(bind, signing_key, server_tp.into(), handler, driver)
    }

    pub fn build_server_with_config(
        bind: SocketAddr,
        signing_key: SigningKey,
        mut server_config: ConnConfig,
        handler: H,
        driver: &mut Driver,
    ) -> io::Result<Self> {
        let udp = Socket::bind(bind, driver)?;
        if server_config.accept_early_data && server_config.early_data_store.is_none() {
            server_config.early_data_store = Some(crate::early_data::shared_replay_store());
        }
        let mux = Mux::server(handler, signing_key, server_config);
        Ok(Self { udp, mux })
    }

    pub fn build_client(bind: SocketAddr, handler: H, driver: &mut Driver) -> io::Result<Self> {
        let udp = Socket::bind(bind, driver)?;
        let mux = Mux::client(handler);
        Ok(Self { udp, mux })
    }

    pub fn local_addr(&self) -> SocketAddr {
        self.udp.local_addr()
    }

    pub fn set_gso(self: Pin<&mut Self>, on: bool) {
        self.project().mux.set_gso(on);
    }

    pub fn handler(&self) -> &H {
        self.mux.handler()
    }

    pub fn handler_mut(self: Pin<&mut Self>) -> &mut H {
        self.project().mux.handler_mut()
    }

    pub fn connect(
        self: Pin<&mut Self>,
        peer_addr: SocketAddr,
        server_pubkey: [u8; 32],
        client_tp: transport_params::Params,
        initial_dcid: Vec<u8>,
    ) -> ConnHandle {
        self.connect_with_config(peer_addr, server_pubkey, client_tp.into(), initial_dcid)
    }

    pub fn connect_with_config(
        self: Pin<&mut Self>,
        peer_addr: SocketAddr,
        server_pubkey: [u8; 32],
        client_config: ConnConfig,
        initial_dcid: Vec<u8>,
    ) -> ConnHandle {
        let now = Instant::now();
        self.project()
            .mux
            .connect(peer_addr, server_pubkey, client_config, initial_dcid, now)
    }

    pub fn conn_mut(self: Pin<&mut Self>, handle: ConnHandle) -> Option<&mut Conn> {
        self.project().mux.conn_mut(handle)
    }

    pub fn send_datagram(
        self: Pin<&mut Self>,
        handle: ConnHandle,
        data: Vec<u8>,
    ) -> Result<(), DatagramError> {
        let now = Instant::now();
        self.project().mux.send_datagram(handle, data, now)
    }

    pub fn close(self: Pin<&mut Self>, handle: ConnHandle) {
        self.project().mux.close(handle);
    }

    pub fn drive(mut self: Pin<&mut Self>, driver: &mut Driver) {
        let mut buf = [Cqe::ZERO; 64];
        loop {
            let n = driver.drain(&mut buf);
            if n == 0 {
                break;
            }
            for cqe in &buf[..n] {
                let Ok(ev) = dope::Event::try_from(*cqe) else {
                    continue;
                };
                self.as_mut().dispatch(ev, driver);
            }
        }
        self.pre_park(driver);
    }
}

impl<const ID: u8, H: mux::Handler> Manifold for Endpoint<ID, H> {
    const ID: u8 = ID;

    fn dispatch(self: Pin<&mut Self>, ev: dope::Event, driver: &mut Driver) {
        use dope::Event;
        let this = self.project();
        match ev {
            Event::Recv(token, more, e) => {
                this.udp.dispatch_recv(token, more, e, driver, this.mux);
            }
            Event::Send(token, e) => {
                this.udp.dispatch_send(token, e, this.mux);
            }
            _ => {}
        }
    }

    fn pre_park(self: Pin<&mut Self>, driver: &mut Driver) {
        let mut this = self.project();
        let now = Instant::now();
        this.mux.reap_closed(now);
        let gso = this.mux.gso();
        flush_coalesced(this.udp.as_mut(), this.mux.pull_outgoing(), gso);
        this.udp.tick(driver);
    }

    fn idle(self: Pin<&Self>) -> dope::Idle {
        if self.get_ref().udp.has_pending() {
            dope::Idle::Busy
        } else {
            dope::Idle::Park(None)
        }
    }
}

impl<H: mux::Handler> Handler<0> for Mux<H> {
    fn on_packet(&mut self, addr: SocketAddr, data: &[u8], mut sock: Pin<&mut Socket<0>>) {
        let now = Instant::now();
        let _ = self.on_udp_packet(addr, data, now);
        let gso = self.gso();
        flush_coalesced(sock.as_mut(), self.pull_outgoing(), gso);
    }
}

const MAX_GSO_BYTES: usize = 65535;
const MAX_GSO_SEGMENTS: usize = 64;

/// Coalesces same-dst runs (equal segments, ≤1 short tail) into one GSO send.
fn flush_coalesced(
    mut sock: Pin<&mut Socket<0>>,
    mut items: Vec<(SocketAddr, Vec<u8>)>,
    gso: bool,
) {
    if !gso {
        for (addr, payload) in items {
            sock.as_mut().queue_to(payload, addr);
        }
        return;
    }
    let mut i = 0;
    while i < items.len() {
        let addr = items[i].0;
        let seg = items[i].1.len();
        let mut j = i + 1;
        let mut total = seg;
        while j < items.len()
            && items[j].0 == addr
            && items[j - 1].1.len() == seg
            && items[j].1.len() <= seg
            && total + items[j].1.len() <= MAX_GSO_BYTES
            && j - i < MAX_GSO_SEGMENTS
        {
            total += items[j].1.len();
            j += 1;
        }
        if j - i == 1 {
            sock.as_mut()
                .queue_to(std::mem::take(&mut items[i].1), addr);
        } else {
            let mut buf = Vec::with_capacity(total);
            for slot in &mut items[i..j] {
                buf.append(&mut slot.1);
            }
            debug_assert!(seg <= u16::MAX as usize);
            sock.as_mut().queue_segmented(buf, addr, seg as u16);
        }
        i = j;
    }
}
