use core::{
    mem::size_of,
    net::{IpAddr, SocketAddr},
    task::{Context, Poll},
};

use alloc::{vec, vec::Vec};
use bytemuck::Zeroable;
use libk_util::waker::QueueWaker;
use yggdrasil_abi::{
    error::Error,
    net::{
        protocols::{InetChecksum, IpProtocol, TcpFlags, TcpFrame, TcpV4PseudoHeader},
        types::NetValueImpl,
    },
};

use crate::{
    l3::{self, L3Packet},
    socket::tcp::{TcpListener, TcpStream},
    util::Assembler,
};

#[derive(PartialEq, Debug)]
pub enum TcpConnectionState {
    SynSent,
    SynReceived,
    Established,
    FinWait1,
    FinWait2,
    Closed,
}

pub enum TcpSocketBehavior {
    None,
    Accept,
    Remove,
}

struct SocketBuffer {
    data: Vec<u8>,
    wr: usize,
    rd: usize,
}

pub struct TcpConnection {
    state: TcpConnectionState,

    local: SocketAddr,
    remote: SocketAddr,

    rx_window_size: u16,

    // Rx half
    // RCV.WND = rx_buffer.capacity()
    // RCV.NXT = rx_window_start + rx_window.len()
    // TODO RCV.UP
    rx_buffer: SocketBuffer,
    rx_segment_buffer: Vec<u8>,
    rx_assembler: Assembler,

    // Relative RX sequence number of window start
    rx_window_start: usize,

    // Tx half
    // SND.UNA = tx_window_start + tx_sent_unacknowledged
    // SND.WND = tx_buffer.capacity()
    // SND.NXT = tx_window_start IF tx_sent_unacknowledged == 0
    tx_buffer: Vec<u8>,

    tx_window_start: usize,
    tx_sent_unacknowledged: usize,

    // IRS
    initial_rx_seq: u32,
    initial_tx_seq: u32,

    rx_notify: QueueWaker,
    tx_notify: QueueWaker,
}

#[allow(unused)]
struct TcpPacket {
    local: SocketAddr,
    remote: SocketAddr,
    seq: u32,
    ack: u32,
    window_size: u16,
    flags: TcpFlags,
}

impl SocketBuffer {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            data: vec![0; capacity],
            wr: 0,
            rd: 0,
        }
    }

    pub fn len(&self) -> usize {
        if self.wr >= self.rd {
            self.wr - self.rd
        } else {
            self.wr + self.capacity() - self.rd
        }
    }

    pub fn capacity(&self) -> usize {
        self.data.len()
    }

    pub fn can_read(&self) -> bool {
        self.rd != self.wr
    }

    pub fn write(&mut self, data: &[u8]) {
        for &byte in data {
            self.putc(byte);
        }
    }

    pub fn putc(&mut self, data: u8) {
        if (self.wr + 1) % self.capacity() == self.rd {
            self.rd = (self.rd + 1) % self.capacity();
        }
        self.data[self.wr] = data;
        self.wr = (self.wr + 1) % self.capacity();
    }

    pub fn read(&mut self, buffer: &mut [u8]) -> usize {
        let mut amount = 0;
        while amount != buffer.len() {
            if self.rd == self.wr {
                break;
            }
            buffer[amount] = self.data[self.rd];
            self.rd = (self.rd + 1) % self.capacity();
            amount += 1;
        }
        amount
    }
}

impl TcpConnection {
    pub fn new(
        local: SocketAddr,
        remote: SocketAddr,
        window_size: usize,
        tx_seq: u32,
        rx_seq: u32,
        state: TcpConnectionState,
    ) -> Self {
        debug_assert!(
            state == TcpConnectionState::SynSent
                || state == TcpConnectionState::Closed
                || state == TcpConnectionState::SynReceived
        );
        debug_assert!(window_size < u16::MAX as usize);
        Self {
            state,

            local,
            remote,

            rx_buffer: SocketBuffer::with_capacity(window_size),
            rx_assembler: Assembler::new(),
            rx_segment_buffer: Vec::with_capacity(window_size),
            rx_window_start: 1,

            tx_buffer: Vec::with_capacity(window_size),
            tx_window_start: 1,
            tx_sent_unacknowledged: 0,

            rx_window_size: window_size as u16,

            initial_rx_seq: rx_seq,
            initial_tx_seq: tx_seq,

            rx_notify: QueueWaker::new(),
            tx_notify: QueueWaker::new(),
        }
    }

    fn ack_number(&self) -> u32 {
        (self.initial_rx_seq as usize + self.rx_window_start + self.rx_buffer.len()) as u32
    }

    fn seq_number(&self) -> u32 {
        (self.initial_tx_seq as usize + self.tx_window_start) as u32
    }

    pub fn is_closing(&self) -> bool {
        self.state == TcpConnectionState::FinWait1
            || self.state == TcpConnectionState::FinWait2
            || self.state == TcpConnectionState::Closed
    }

    pub fn is_closed(&self) -> bool {
        self.state == TcpConnectionState::Closed
    }

    pub fn read_nonblocking(&mut self, buffer: &mut [u8]) -> Result<usize, Error> {
        let amount = self.rx_buffer.read(buffer);
        if amount == 0 && self.state != TcpConnectionState::Established && !self.is_closing() {
            // TODO ConnectionAborted?
            return Err(Error::ConnectionReset);
        }
        self.rx_window_start += amount;
        Ok(amount)
    }

    pub(crate) fn poll_receive(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.rx_notify.register(cx.waker());
        if self.rx_buffer.can_read() {
            self.rx_notify.remove(cx.waker());
            Poll::Ready(Ok(()))
        } else if self.state != TcpConnectionState::Established {
            self.rx_notify.remove(cx.waker());
            Poll::Ready(Err(Error::ConnectionReset))
        } else {
            Poll::Pending
        }
    }

    pub(crate) fn poll_send(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.tx_notify.register(cx.waker());
        if self.state == TcpConnectionState::Closed {
            self.tx_notify.remove(cx.waker());
            Poll::Ready(Err(Error::ConnectionReset))
        } else if self.tx_sent_unacknowledged == 0 {
            self.tx_notify.remove(cx.waker());
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    pub(crate) fn poll_acknowledge(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.tx_notify.register(cx.waker());
        if self.tx_sent_unacknowledged == 0 {
            self.tx_notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    pub(crate) fn poll_finish(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.tx_notify.register(cx.waker());
        if self.state == TcpConnectionState::FinWait2 || self.state == TcpConnectionState::Closed {
            self.tx_notify.remove(cx.waker());
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    pub(crate) fn poll_established(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.rx_notify.register(cx.waker());
        match self.state {
            TcpConnectionState::Established => Poll::Ready(Ok(())),
            TcpConnectionState::Closed => Poll::Ready(Err(Error::ConnectionRefused)),
            _ => Poll::Pending,
        }
    }

    pub(crate) async fn transmit(&mut self, data: &[u8]) -> Result<(), Error> {
        assert_eq!(self.tx_sent_unacknowledged, 0);
        assert_eq!(self.tx_buffer.len(), 0);

        self.tx_buffer.extend_from_slice(data);
        self.tx_sent_unacknowledged = data.len();

        send(
            self.local,
            self.remote,
            self.seq_number(),
            self.ack_number(),
            self.rx_window_size,
            TcpFlags::ACK,
            data,
        )
        .await
    }

    pub(crate) async fn finish(&mut self) -> Result<(), Error> {
        assert_eq!(self.tx_sent_unacknowledged, 0);
        assert_eq!(self.tx_buffer.len(), 0);

        log::debug!("Finish connection {} <-> {}", self.local, self.remote);
        send(
            self.local,
            self.remote,
            self.seq_number(),
            self.ack_number(),
            self.rx_window_size,
            TcpFlags::FIN | TcpFlags::ACK,
            &[],
        )
        .await?;

        self.state = TcpConnectionState::FinWait1;

        Ok(())
    }

    pub(crate) fn notify_all(&self) {
        self.rx_notify.wake_all();
        self.tx_notify.wake_all();
    }

    pub(crate) async fn send_syn(&mut self) -> Result<(), Error> {
        assert!(
            self.state == TcpConnectionState::SynSent || self.state == TcpConnectionState::Closed
        );
        log::debug!("Send SYN {} -> {}", self.local, self.remote);
        self.state = TcpConnectionState::SynSent;
        send(
            self.local,
            self.remote,
            self.initial_tx_seq,
            0,
            self.rx_window_size,
            TcpFlags::SYN,
            &[],
        )
        .await
    }

    fn handle_packet_payload(&mut self, data: &[u8], seq: u32) -> bool {
        // Local side
        let rx_window_start = self.rx_window_start + self.rx_buffer.len();
        let rx_window_end = self.rx_window_start + self.rx_buffer.capacity();
        // Remote side
        let rx_segment_start = seq.wrapping_sub(self.initial_rx_seq) as usize;
        let rx_segment_end = rx_segment_start + data.len();

        if rx_segment_end >= rx_window_end || rx_segment_start < rx_window_start {
            return false;
        }

        // Offset from expected seq
        let segment_start_offset = rx_segment_start - rx_window_start;

        // Push the data into reassembler buffer
        assert!(segment_start_offset + data.len() <= self.rx_segment_buffer.capacity());
        if segment_start_offset + data.len() > self.rx_segment_buffer.len() {
            self.rx_segment_buffer
                .resize(segment_start_offset + data.len(), 0);
        }
        self.rx_segment_buffer[segment_start_offset..segment_start_offset + data.len()]
            .copy_from_slice(data);

        let amount = self
            .rx_assembler
            .add_then_remove_front(segment_start_offset, data.len())
            .unwrap();

        if amount != 0 {
            // Take data from reassembly buffer and append to rx_buffer
            self.rx_buffer.write(&self.rx_segment_buffer[..amount]);
            self.rx_segment_buffer.drain(..amount);
            self.rx_notify.wake_one();
        }

        true
    }

    async fn handle_packet(
        &mut self,
        packet: TcpPacket,
        data: &[u8],
    ) -> Result<TcpSocketBehavior, Error> {
        // TODO what if window_size changes?

        match self.state {
            TcpConnectionState::SynSent => {
                if packet.flags == TcpFlags::SYN | TcpFlags::ACK {
                    if packet.ack != self.initial_tx_seq.wrapping_add(1) {
                        log::warn!(
                            "Expected ACK {}, got {}",
                            self.initial_tx_seq.wrapping_add(1),
                            packet.ack
                        );
                        return Ok(TcpSocketBehavior::None);
                    }

                    log::debug!(
                        "TCP {} -> {} got ACKed, established",
                        self.local,
                        self.remote
                    );
                    self.initial_rx_seq = packet.seq;

                    // ACK the SYN+ACK
                    send(
                        self.local,
                        self.remote,
                        self.initial_tx_seq.wrapping_add(1),
                        self.initial_rx_seq.wrapping_add(1),
                        self.rx_window_size,
                        TcpFlags::ACK,
                        &[],
                    )
                    .await?;

                    self.state = TcpConnectionState::Established;
                    self.rx_notify.wake_all();
                } else if packet.flags == TcpFlags::RST | TcpFlags::ACK {
                    log::debug!("TCP {} -> {} got RSTd, closing", self.local, self.remote);
                    self.state = TcpConnectionState::Closed;
                    self.rx_notify.wake_all();
                    return Ok(TcpSocketBehavior::Remove);
                }

                // TODO try re-sending SYN?
                return Ok(TcpSocketBehavior::None);
            }
            TcpConnectionState::FinWait1 => {
                // Check if connection close initiated locally got ACKed by remote
                // TODO check ack/seq
                if packet.flags == TcpFlags::FIN | TcpFlags::ACK {
                    self.state = TcpConnectionState::Closed;
                    send(
                        self.local,
                        self.remote,
                        self.seq_number() + 1,
                        self.ack_number() + 1,
                        self.rx_window_size,
                        TcpFlags::ACK,
                        &[],
                    )
                    .await?;

                    // Socket fully closed, remove from table
                    return Ok(TcpSocketBehavior::Remove);
                }
                if packet.flags == TcpFlags::ACK {
                    self.state = TcpConnectionState::FinWait2;
                    self.tx_notify.wake_all();
                }
                return Ok(TcpSocketBehavior::None);
            }
            TcpConnectionState::FinWait2 => {
                if packet.flags == TcpFlags::FIN | TcpFlags::ACK {
                    self.state = TcpConnectionState::Closed;
                    send(
                        self.local,
                        self.remote,
                        self.seq_number() + 1,
                        self.ack_number() + 1,
                        self.rx_window_size,
                        TcpFlags::ACK,
                        &[],
                    )
                    .await?;

                    // Socket fully closed, remove from table
                    return Ok(TcpSocketBehavior::Remove);
                }

                return Ok(TcpSocketBehavior::None);
            }
            TcpConnectionState::Closed => {
                log::warn!("Packet received on closed connection");
                return Ok(TcpSocketBehavior::None);
            }
            TcpConnectionState::SynReceived => {
                // TODO check ack/seq
                // Handshake continuation expected (ACK)
                if packet.flags == TcpFlags::ACK {
                    self.state = TcpConnectionState::Established;
                    return Ok(TcpSocketBehavior::Accept);
                }

                return Err(Error::InvalidArgument);
            }
            TcpConnectionState::Established => (),
        }

        if self.tx_sent_unacknowledged != 0 {
            let tx_acknowledge_end = packet.ack.wrapping_sub(self.initial_tx_seq) as usize;

            if tx_acknowledge_end == self.tx_window_start + self.tx_sent_unacknowledged {
                self.tx_window_start += self.tx_sent_unacknowledged;
                self.tx_sent_unacknowledged = 0;
                self.tx_buffer.clear();
                self.tx_notify.wake_one();
            }
        }

        let mut reply_flags = TcpFlags::empty();
        let mut behavior = TcpSocketBehavior::None;

        if !data.is_empty() && self.handle_packet_payload(data, packet.seq) {
            reply_flags |= TcpFlags::ACK;
        }

        // TODO check window resize notification

        let mut ack_number = self.ack_number();

        if packet.flags.contains(TcpFlags::FIN) {
            reply_flags |= TcpFlags::ACK;
            ack_number = ack_number.wrapping_add(1);

            // Only send an actual FIN after a FIN without any data
            if data.is_empty() {
                reply_flags |= TcpFlags::FIN;
                // TODO go to LastAck state and wait for ACK
                self.state = TcpConnectionState::Closed;
                log::trace!(
                    "TCP connection FIN requested by remote: {} <-> {}",
                    self.local,
                    self.remote
                );

                behavior = TcpSocketBehavior::Remove;
            }
        }

        if reply_flags != TcpFlags::empty() {
            send(
                packet.local,
                packet.remote,
                self.seq_number(),
                ack_number,
                self.rx_window_size,
                reply_flags,
                &[],
            )
            .await?;
        }

        Ok(behavior)
    }
}

async fn send(
    local: SocketAddr,
    remote: SocketAddr,
    seq: u32,
    ack: u32,
    window_size: u16,
    flags: TcpFlags,
    data: &[u8],
) -> Result<(), Error> {
    let (interface, source_ip, gateway_ip, destination_ip) = l3::resolve_l3_route(remote.ip())?;
    // TODO TCPv6
    let source_ip = match source_ip {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };
    let destination_ip = match destination_ip {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };

    let tcp_length = size_of::<TcpFrame>() + data.len();

    let mut frame = TcpFrame {
        source_port: local.port().to_network_order(),
        destination_port: remote.port().to_network_order(),
        sequence_number: seq.to_network_order(),
        acknowledge_number: ack.to_network_order(),
        data_offset: 5 << 4,
        window_size: window_size.to_network_order(),
        flags,
        ..TcpFrame::zeroed()
    };
    let pseudo_header = TcpV4PseudoHeader {
        source_address: u32::from(source_ip).to_network_order(),
        destination_address: u32::from(destination_ip).to_network_order(),
        _zero: 0,
        protocol: IpProtocol::TCP,
        tcp_length: (tcp_length as u16).to_network_order(),
    };

    let mut checksum = InetChecksum::new();
    checksum.add_value(&pseudo_header, true);
    checksum.add_value(&frame, true);
    checksum.add_bytes(data, true);
    let checksum = checksum.finish();

    frame.checksum = checksum.to_network_order();

    l3::send_l4_ip_resolved(&l3::L4ResolvedPacket {
        interface: &interface,
        source_ip: source_ip.into(),
        gateway_ip,
        destination_ip: destination_ip.into(),
        l4_frame: bytemuck::bytes_of(&frame),
        l4_options: &[],
        l4_data: data,
        protocol: IpProtocol::TCP,
        ttl: 64,
    })
    .await
}

fn validate(source: IpAddr, destination: IpAddr, tcp_frame: &TcpFrame, data: &[u8]) -> bool {
    // TODO TCPv6
    let source = match source {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };
    let destination = match destination {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };
    let tcp_length = size_of::<TcpFrame>() + data.len();

    let pseudo_header = TcpV4PseudoHeader {
        source_address: u32::from(source).to_network_order(),
        destination_address: u32::from(destination).to_network_order(),
        _zero: 0,
        protocol: IpProtocol::TCP,
        tcp_length: (tcp_length as u16).to_network_order(),
    };

    let mut checksum = InetChecksum::new();
    checksum.add_value(&pseudo_header, true);
    checksum.add_value(tcp_frame, true);
    checksum.add_bytes(data, true);
    let checksum = checksum.finish();

    checksum == 0
}

pub async fn handle(packet: L3Packet) -> Result<(), Error> {
    if packet.data_length < size_of::<TcpFrame>() {
        log::warn!("Truncated TCP packet");
        return Ok(());
    }

    let l3_data = packet.l3_data();

    let tcp_frame: &TcpFrame = bytemuck::from_bytes(&l3_data[..size_of::<TcpFrame>()]);
    let tcp_data_offset = tcp_frame.data_offset();
    let tcp_data = &l3_data[tcp_data_offset..packet.data_length];

    let remote = SocketAddr::new(
        packet.source_address,
        u16::from_network_order(tcp_frame.source_port),
    );
    let local = SocketAddr::new(
        packet.destination_address,
        u16::from_network_order(tcp_frame.destination_port),
    );

    let seq = u32::from_network_order(tcp_frame.sequence_number);
    let ack = u32::from_network_order(tcp_frame.acknowledge_number);

    if !validate(
        packet.source_address,
        packet.destination_address,
        tcp_frame,
        &l3_data[size_of::<TcpFrame>()..packet.data_length],
    ) {
        log::warn!("Invalid TCP packet received");
        return Ok(());
    }

    match tcp_frame.flags {
        TcpFlags::SYN => {
            if let Some(listener) = TcpListener::get(&local)
                && listener.is_listening()
            {
                log::debug!("tcp: create remote stream {remote} -> {local}");
                let window_size = u16::from_network_order(tcp_frame.window_size);
                let tx_seq = 12345;

                // Create a socket and insert it into the table
                TcpStream::accept_remote(
                    listener.clone(),
                    local,
                    remote,
                    window_size as usize,
                    tx_seq,
                    seq,
                )?;

                // Send SYN+ACK
                send(
                    local,
                    remote,
                    tx_seq,
                    seq.wrapping_add(1),
                    window_size,
                    TcpFlags::SYN | TcpFlags::ACK,
                    &[],
                )
                .await
            } else {
                // RST+ACK
                log::warn!("SYN {} -> {}: port not listening", remote, local);
                let window_size = u16::from_network_order(tcp_frame.window_size);
                send(
                    local,
                    remote,
                    0,
                    seq.wrapping_add(1),
                    window_size,
                    TcpFlags::RST | TcpFlags::ACK,
                    &[],
                )
                .await
            }
        }
        _ => {
            let packet = TcpPacket {
                local,
                remote,
                window_size: u16::from_network_order(tcp_frame.window_size),
                flags: tcp_frame.flags,
                ack,
                seq,
            };

            let stream = TcpStream::get(local, remote).ok_or(Error::DoesNotExist)?;
            let mut connection = stream.connection.write();

            match connection.handle_packet(packet, tcp_data).await? {
                TcpSocketBehavior::None => (),
                TcpSocketBehavior::Accept => {
                    stream.accept();
                }
                TcpSocketBehavior::Remove => {
                    drop(connection);
                    stream.remove_stream()?;
                }
            }
            Ok(())
        }
    }
}
