use std::error;
use std::fmt::{Debug, Display};
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;

use cyphernet::encrypt::noise::NoiseState;
use cyphernet::proxy::socks5;

use mio::event::Source;
use mio::net::TcpStream;
use mio::{Interest, Registry, Token};

pub type NoiseSession<E, D, S> = Protocol<NoiseState<E, D>, S>;
pub type Socks5Session<S> = Protocol<socks5::Socks5, S>;

pub trait Session: Send + Read + Write {
    type Artifact: Display;

    fn is_established(&self) -> bool {
        self.artifact().is_some()
    }

    fn artifact(&self) -> Option<Self::Artifact>;
}

pub trait StateMachine: Sized + Send {
    const NAME: &'static str;

    type Artifact;

    type Error: error::Error + Send + Sync + 'static;

    fn next_read_len(&self) -> usize;

    fn advance(&mut self, input: &[u8]) -> Result<Vec<u8>, Self::Error>;

    fn artifact(&self) -> Option<Self::Artifact>;

    fn is_complete(&self) -> bool {
        self.artifact().is_some()
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
pub struct ProtocolArtifact<M: StateMachine, S: Session> {
    pub(crate) session: S::Artifact,
    pub(crate) state: M::Artifact,
}

impl<M: StateMachine, S: Session> Display for ProtocolArtifact<M, S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProtocolArtifact")
            .field("session", &"<omitted>")
            .field("state", &"<omitted>")
            .finish()
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub struct Protocol<M: StateMachine, S: Session> {
    pub(crate) machine: M,
    pub(crate) session: S,
}

impl<M: StateMachine, S: Session> Protocol<M, S> {
    pub fn new(session: S, machine: M) -> Self {
        Self { machine, session }
    }
}

impl<M: StateMachine, S: Session> io::Read for Protocol<M, S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        log::trace!(target: M::NAME, "Reading event");

        if self.machine.is_complete() || !self.session.is_established() {
            log::trace!(target: M::NAME, "Passing reading to inner not yet established session");
            return self.session.read(buf);
        }

        let len = self.machine.next_read_len();
        let mut input = vec![0u8; len];
        self.session.read_exact(&mut input)?;

        log::trace!(target: M::NAME, "Received handshake act: {input:02x?}");

        if !input.is_empty() {
            let output = self.machine.advance(&input).map_err(|err| {
                log::error!(target: M::NAME, "Handshake failure: {err}");
                io::Error::other(err)
            })?;

            if !output.is_empty() {
                log::trace!(target: M::NAME, "Sending handshake act on read: {output:02x?}");
                self.session.write_all(&output)?;
            }
        }

        Ok(0)
    }
}

impl<M: StateMachine, S: Session> Write for Protocol<M, S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        log::trace!(target: M::NAME, "Writing event (state_complete={}, session_established={})", self.machine.is_complete(), self.session.is_established());

        if self.machine.is_complete() || !self.session.is_established() {
            log::trace!(target: M::NAME, "Passing writing to inner session");
            return self.session.write(buf);
        }

        if self.machine.next_read_len() == 0 {
            log::trace!(target: M::NAME, "Starting handshake protocol");

            let act = self.machine.advance(&[]).map_err(|err| {
                log::error!(target: M::NAME, "Handshake failure: {err}");
                io::Error::other(err)
            })?;

            if !act.is_empty() {
                log::trace!(target: M::NAME, "Sending handshake act on write: {act:02x?}");
                self.session.write_all(&act)?;
            } else {
                log::trace!(target: M::NAME, "Handshake complete, passing data to inner session");
                return self.session.write(buf);
            }
        }

        if buf.is_empty() {
            Ok(0)
        } else {
            Err(io::ErrorKind::Interrupted.into())
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        self.session.flush()
    }
}

impl<M: StateMachine, S: Session> Session for Protocol<M, S> {
    type Artifact = ProtocolArtifact<M, S>;

    fn artifact(&self) -> Option<Self::Artifact> {
        Some(ProtocolArtifact {
            session: self.session.artifact()?,
            state: self.machine.artifact()?,
        })
    }
}

impl<M: StateMachine, S: Session + Source> Source for Protocol<M, S> {
    fn register(
        &mut self,
        registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.session.register(registry, token, interests)
    }

    fn reregister(
        &mut self,
        registry: &Registry,
        token: Token,
        interests: Interest,
    ) -> io::Result<()> {
        self.session.reregister(registry, token, interests)
    }

    fn deregister(&mut self, registry: &Registry) -> io::Result<()> {
        self.session.deregister(registry)
    }
}

impl Session for TcpStream {
    type Artifact = SocketAddr;

    fn artifact(&self) -> Option<Self::Artifact> {
        self.peer_addr().ok()
    }
}

mod impl_noise {
    use cyphernet::encrypt::noise::{NoiseState as Noise, error::NoiseError as Error};
    use cyphernet::{Digest, Ecdh};

    use super::*;

    #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
    pub struct NoiseArtifact<E: Ecdh, D: Digest> {
        pub handshake_hash: D::Output,
        pub remote_static_key: Option<E::Pk>,
    }

    impl<E: Ecdh, D: Digest> StateMachine for Noise<E, D> {
        const NAME: &'static str = "noise";
        type Artifact = NoiseArtifact<E, D>;
        type Error = Error;

        fn next_read_len(&self) -> usize {
            self.next_read_len()
        }

        fn advance(&mut self, input: &[u8]) -> Result<Vec<u8>, Self::Error> {
            self.advance(input)
        }

        fn artifact(&self) -> Option<Self::Artifact> {
            self.get_handshake_hash().map(|hh| NoiseArtifact {
                handshake_hash: hh,
                remote_static_key: self.get_remote_static_key(),
            })
        }
    }
}

mod impl_socks5 {
    use cyphernet::addr::{Host as _, HostName, NetAddr};
    use cyphernet::proxy::socks5::{Error, Socks5};

    use super::*;

    impl StateMachine for Socks5 {
        const NAME: &'static str = "socks5";

        type Artifact = NetAddr<HostName>;
        type Error = Error;

        fn next_read_len(&self) -> usize {
            self.next_read_len()
        }

        fn advance(&mut self, input: &[u8]) -> Result<Vec<u8>, Self::Error> {
            self.advance(input)
        }

        fn artifact(&self) -> Option<Self::Artifact> {
            match self {
                Socks5::Initial(addr, false) if !addr.requires_proxy() => Some(addr.clone()),
                Socks5::Active(addr) => Some(addr.clone()),
                _ => None,
            }
        }
    }
}
