#![allow(dead_code)]

use asynchronous_codec::Framed;
use bytes::BytesMut;
use futures::prelude::*;
use libp2p::core::{upgrade, InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use unsigned_varint::codec::UviBytes;

const LOG_TARGET: &str = "ghost-walker-upgrades";

use std::{
    convert::Infallible,
    io, mem,
    pin::Pin,
    task::{Context, Poll},
    vec,
};

const MAX_HANDSHAKE_SIZE: usize = 1024;

#[derive(Debug, Clone)]
pub struct HandshakeInbound {
    pub name: String,
}

#[pin_project::pin_project]
pub struct HandshakeInboundSubstream<TSubstream> {
    #[pin]
    socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
    state: HandshakeInboundSubstreamState,
    negotiated_name: String,
}

#[derive(Debug)]
pub enum HandshakeInboundSubstreamState {
    Waiting,
    Sending(Vec<u8>),
    Flush,
    Done,
    NeedsClose,
    FullyClosed,
}

#[derive(Debug, Clone)]
pub struct HandshakeOutbound {
    pub name: String,
    pub handshake: Vec<u8>,
}

#[pin_project::pin_project]
pub struct HandshakeOutboundSubstream<TSubstream> {
    #[pin]
    socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
}

impl HandshakeInbound {
    pub fn new(name: impl Into<String>) -> Self {
        Self { name: name.into() }
    }
}

impl UpgradeInfo for HandshakeInbound {
    type Info = String;
    type InfoIter = vec::IntoIter<Self::Info>;

    fn protocol_info(&self) -> Self::InfoIter {
        vec![self.name.clone()].into_iter()
    }
}

impl<TSubstream> InboundUpgrade<TSubstream> for HandshakeInbound
where
    TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    type Output = HandshakeInboundOpen<TSubstream>;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
    type Error = HandshakeError;

    fn upgrade_inbound(self, mut socket: TSubstream, negotiated_name: Self::Info) -> Self::Future {
        tracing::debug!(
            target: LOG_TARGET,
            "HandshakeInbound name={:?} current_name={:?}",
            negotiated_name,
            self.name
        );

        Box::pin(async move {
            let handshake_len = unsigned_varint::aio::read_usize(&mut socket).await?;
            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeInbound length={:?} name={:?}",
                handshake_len,
                negotiated_name
            );

            if handshake_len > MAX_HANDSHAKE_SIZE {
                return Err(HandshakeError::TooLarge {
                    requested: handshake_len,
                    max: MAX_HANDSHAKE_SIZE,
                });
            }

            let mut handshake = vec![0u8; handshake_len];
            if !handshake.is_empty() {
                socket.read_exact(&mut handshake).await?;
            }

            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeInbound received handshake={:?} name={:?}",
                handshake,
                negotiated_name
            );

            let mut codec: UviBytes<io::Cursor<Vec<u8>>> = UviBytes::default();
            codec.set_max_len(usize::MAX);

            let substream = HandshakeInboundSubstream {
                socket: Framed::new(socket, codec),
                state: HandshakeInboundSubstreamState::Waiting,
                negotiated_name,
            };

            Ok(HandshakeInboundOpen {
                handshake,
                substream,
            })
        })
    }
}

pub struct HandshakeInboundOpen<TSubstream> {
    pub handshake: Vec<u8>,
    pub substream: HandshakeInboundSubstream<TSubstream>,
}

impl<TSubstream> HandshakeInboundSubstream<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite + Unpin,
{
    pub fn set_handshake(&mut self, handshake: Vec<u8>) {
        match &self.state {
            HandshakeInboundSubstreamState::Waiting => (),
            _ => return,
        };

        self.state = HandshakeInboundSubstreamState::Sending(handshake);
    }

    pub fn poll_process(
        self: Pin<&mut Self>,
        cx: &mut Context,
    ) -> Poll<Result<Infallible, io::Error>> {
        let mut this = self.project();

        loop {
            let state = mem::replace(this.state, HandshakeInboundSubstreamState::Done);

            match state {
                HandshakeInboundSubstreamState::Sending(handshake) => {
                    match Sink::poll_ready(this.socket.as_mut(), cx) {
                        Poll::Ready(_) => {
                            tracing::trace!(target: LOG_TARGET, "HandshakeInboundSubstream: poll_process: Sink is ready start sending name={:?}", this.negotiated_name);

                            *this.state = HandshakeInboundSubstreamState::Flush;

                            match Sink::start_send(this.socket.as_mut(), io::Cursor::new(handshake))
                            {
                                Ok(()) => {}
                                Err(err) => {
                                    tracing::error!(target: LOG_TARGET, "HandshakeInboundSubstream: poll_process: Failed to start sending name={:?} error={:?}", this.negotiated_name, err);

                                    return Poll::Ready(Err(err));
                                }
                            }
                        }
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::Sending(handshake);
                            return Poll::Pending;
                        }
                    }
                }

                HandshakeInboundSubstreamState::Flush => {
                    tracing::trace!(
                        target: LOG_TARGET,
                        "HandshakeInboundSubstream: poll_process: poll_flush name={:?}",
                        this.negotiated_name
                    );
                    match Sink::poll_flush(this.socket.as_mut(), cx)? {
                        Poll::Ready(()) => *this.state = HandshakeInboundSubstreamState::Done,
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::Flush;
                            return Poll::Pending;
                        }
                    }
                }

                state => {
                    *this.state = state;
                    return Poll::Pending;
                }
            }
        }
    }
}

impl<TSubstream> Stream for HandshakeInboundSubstream<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite + Unpin,
{
    type Item = Result<BytesMut, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        loop {
            let state = mem::replace(this.state, HandshakeInboundSubstreamState::Done);

            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeInboundSubstream: poll_next: state={:?} name={:?}",
                state,
                this.negotiated_name
            );

            match state {
                HandshakeInboundSubstreamState::Waiting => {
                    *this.state = HandshakeInboundSubstreamState::Waiting;
                    return Poll::Pending;
                }
                HandshakeInboundSubstreamState::Sending(handshake) => {
                    match Sink::poll_ready(this.socket.as_mut(), cx) {
                        Poll::Ready(_) => {
                            *this.state = HandshakeInboundSubstreamState::Flush;

                            match Sink::start_send(this.socket.as_mut(), io::Cursor::new(handshake))
                            {
                                Ok(()) => (),
                                Err(err) => {
                                    tracing::trace!(
                                        target: LOG_TARGET,
                                        "HandshakeInboundSubstream: Cannot send handshake name={:?}",
                                        this.negotiated_name
                                    );

                                    return Poll::Ready(Some(Err(err)));
                                }
                            }
                        }
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::Sending(handshake);
                            return Poll::Pending;
                        }
                    }
                }
                HandshakeInboundSubstreamState::Flush => {
                    match Sink::poll_flush(this.socket.as_mut(), cx)? {
                        Poll::Ready(()) => *this.state = HandshakeInboundSubstreamState::Done,
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::Flush;
                            return Poll::Pending;
                        }
                    }
                }
                HandshakeInboundSubstreamState::Done => {
                    match Stream::poll_next(this.socket.as_mut(), cx) {
                        Poll::Ready(None) => {
                            tracing::trace!(
                                target: LOG_TARGET,
                                "HandshakeInboundSubstream: poll_next: Closing in response to peer name={:?}",
                                this.negotiated_name
                            );
                            *this.state = HandshakeInboundSubstreamState::NeedsClose
                        }
                        Poll::Ready(Some(result)) => {
                            *this.state = HandshakeInboundSubstreamState::Done;
                            return Poll::Ready(Some(result));
                        }
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::Done;
                            return Poll::Pending;
                        }
                    }
                }
                HandshakeInboundSubstreamState::NeedsClose => {
                    match Sink::poll_close(this.socket.as_mut(), cx)? {
                        Poll::Ready(()) => {
                            tracing::trace!(
                                target: LOG_TARGET,
                                "HandshakeInboundSubstream: poll_close: fully clsoed name={:?}",
                                this.negotiated_name
                            );
                            *this.state = HandshakeInboundSubstreamState::FullyClosed
                        }
                        Poll::Pending => {
                            *this.state = HandshakeInboundSubstreamState::NeedsClose;
                            return Poll::Pending;
                        }
                    }
                }
                HandshakeInboundSubstreamState::FullyClosed => return Poll::Ready(None),
            }
        }
    }
}

impl HandshakeOutbound {
    pub fn new(name: impl Into<String>, handshake: impl Into<Vec<u8>>) -> Self {
        Self {
            name: name.into(),
            handshake: handshake.into(),
        }
    }
}

impl UpgradeInfo for HandshakeOutbound {
    type Info = String;
    type InfoIter = vec::IntoIter<Self::Info>;

    fn protocol_info(&self) -> Self::InfoIter {
        vec![self.name.clone()].into_iter()
    }
}

impl<TSubstream> OutboundUpgrade<TSubstream> for HandshakeOutbound
where
    TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    type Output = HandshakeOutboundOpen<TSubstream>;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
    type Error = HandshakeError;

    fn upgrade_outbound(self, mut socket: TSubstream, negotiated_name: Self::Info) -> Self::Future {
        tracing::debug!(
            target: LOG_TARGET,
            "HandshakeOutbound name={:?} current_name={:?}",
            negotiated_name,
            self.name
        );

        Box::pin(async move {
            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeOutbound prepare to write handshake={:?} name={:?}",
                self.handshake,
                negotiated_name
            );

            upgrade::write_length_prefixed(&mut socket, &self.handshake).await?;

            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeOutbound prepare to read handshake length name={:?}",
                negotiated_name
            );

            let handshake_len = unsigned_varint::aio::read_usize(&mut socket).await?;

            tracing::trace!(
                target: LOG_TARGET,
                "HandshakeOutbound handshake len={:?} name={:?}",
                handshake_len,
                negotiated_name
            );

            if handshake_len > MAX_HANDSHAKE_SIZE {
                return Err(HandshakeError::TooLarge {
                    requested: handshake_len,
                    max: MAX_HANDSHAKE_SIZE,
                });
            }

            let mut handshake = vec![0u8; handshake_len];
            if !handshake.is_empty() {
                socket.read_exact(&mut handshake).await?;
            }

            let mut codec = UviBytes::default();
            codec.set_max_len(usize::MAX);

            Ok(HandshakeOutboundOpen {
                handshake,
                substream: HandshakeOutboundSubstream {
                    socket: Framed::new(socket, codec),
                },
            })
        })
    }
}

pub struct HandshakeOutboundOpen<TSubstream> {
    pub handshake: Vec<u8>,
    pub substream: HandshakeOutboundSubstream<TSubstream>,
}

impl<TSubstream> Sink<Vec<u8>> for HandshakeOutboundSubstream<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite + Unpin,
{
    type Error = std::io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        Sink::poll_ready(this.socket.as_mut(), cx)
    }

    fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
        let mut this = self.project();
        Sink::start_send(this.socket.as_mut(), io::Cursor::new(item))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        Sink::poll_flush(this.socket.as_mut(), cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        Sink::poll_close(this.socket.as_mut(), cx)
    }
}

#[derive(Debug, thiserror::Error)]
pub enum HandshakeError {
    #[error(transparent)]
    Io(#[from] io::Error),
    #[error("Initial message or handshake was too large: {requested}")]
    TooLarge {
        requested: usize,
        max: usize,
    },
    #[error(transparent)]
    VarintDecode(#[from] unsigned_varint::decode::Error),
}

impl From<unsigned_varint::io::ReadError> for HandshakeError {
    fn from(err: unsigned_varint::io::ReadError) -> Self {
        match err {
            unsigned_varint::io::ReadError::Io(err) => Self::Io(err),
            unsigned_varint::io::ReadError::Decode(err) => Self::VarintDecode(err),
            _ => {
                tracing::warn!(target: LOG_TARGET, "Unrecognized varint decoding error");
                Self::Io(From::from(io::ErrorKind::InvalidData))
            }
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum HandshakeOutboundError {
    #[error(transparent)]
    Io(#[from] io::Error),
}
