use mio::net::{TcpListener, TcpStream, UdpSocket};
use mio::{event, Poll, Interest, Token, Events, Registry};

use std::net::{SocketAddr, SocketAddrV4, Ipv4Addr, TcpStream as StdTcpStream};
use net2::{UdpBuilder};
use std::sync::{Arc, Mutex};
use std::time::{Duration};
use std::collections::{HashMap};
use std::io::{self, prelude::*, ErrorKind};

const EVENTS_SIZE: usize = 1024;
pub const MAX_UDP_LEN: usize = 1488;

/// Information to identify the remote endpoint.
/// The endpoint is used mainly as a connection identified.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct Endpoint {
    resource_id: usize,
    addr: SocketAddr,
}

impl Endpoint {
    fn new(resource_id: usize, addr: SocketAddr) -> Endpoint {
        Endpoint { resource_id, addr }
    }

    /// Returns the inner network resource id associated used for the endpoint.
    /// It is not necessary to be unique for each endpoint, if some of them shared the resource (an example of this is the different endpoints generated by the a UDP Listener).
    pub fn resource_id(&self) -> usize {
        self.resource_id
    }

    /// Returns the remote address of the endpoint
    pub fn addr(&self) -> SocketAddr {
        self.addr
    }
}

impl std::fmt::Display for Endpoint {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "[{}]-{}", self.resource_id, self.addr)
    }
}

#[derive(Debug)]
pub enum Event<'a> {
    Connection,
    Data(&'a [u8]),
    Disconnection,
}

pub enum Listener {
    Tcp(TcpListener),
    Udp(UdpSocket),
}

impl Listener {
    pub fn new_tcp(addr: SocketAddr) -> io::Result<Listener> {
        TcpListener::bind(addr).map(Listener::Tcp)
    }

    pub fn new_udp(addr: SocketAddr) -> io::Result<Listener> {
        UdpSocket::bind(addr).map(Listener::Udp)
    }

    pub fn new_udp_multicast(addr: SocketAddrV4) -> io::Result<Listener> {
        let listening_addr = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, addr.port());
        UdpBuilder::new_v4().unwrap().reuse_address(true).unwrap().bind(listening_addr).map(
            |socket| {
                socket.set_nonblocking(true).unwrap();
                socket.join_multicast_v4(&addr.ip(), &Ipv4Addr::UNSPECIFIED).unwrap();
                Listener::Udp(UdpSocket::from_std(socket))
            },
        )
    }

    pub fn local_addr(&self) -> SocketAddr {
        match self {
            Listener::Tcp(listener) => listener.local_addr().unwrap(),
            Listener::Udp(socket) => socket.local_addr().unwrap(),
        }
    }

    pub fn event_source(&mut self) -> &mut dyn event::Source {
        match self {
            Listener::Tcp(listener) => listener,
            Listener::Udp(socket) => socket,
        }
    }
}

impl Drop for Listener {
    fn drop(&mut self) {
        if let Listener::Udp(socket) = self {
            if let SocketAddr::V4(addr) = socket.local_addr().unwrap() {
                if addr.ip().is_multicast() {
                    socket.leave_multicast_v4(&addr.ip(), &Ipv4Addr::UNSPECIFIED).unwrap();
                }
            }
        }
    }
}

pub enum Remote {
    Tcp(TcpStream),
    Udp(UdpSocket, SocketAddr),
}

impl Remote {
    pub fn new_tcp(addr: SocketAddr) -> io::Result<Remote> {
        // Create a standard TcpStream to blocking until the connection is reached.
        StdTcpStream::connect(addr).map(|stream| {
            stream.set_nonblocking(true).unwrap();
            Remote::Tcp(TcpStream::from_std(stream))
        })
    }

    pub fn new_udp(addr: SocketAddr) -> io::Result<Remote> {
        UdpSocket::bind("0.0.0.0:0".parse().unwrap()).map(|socket| {
            socket.connect(addr).unwrap();
            Remote::Udp(socket, addr)
        })
    }

    pub fn local_addr(&self) -> SocketAddr {
        match self {
            Remote::Tcp(stream) => stream.local_addr().unwrap(),
            Remote::Udp(socket, _) => socket.local_addr().unwrap(),
        }
    }

    pub fn peer_addr(&self) -> SocketAddr {
        match self {
            Remote::Tcp(stream) => stream.peer_addr().unwrap(),
            Remote::Udp(_, addr) => *addr,
        }
    }

    pub fn event_source(&mut self) -> &mut dyn event::Source {
        match self {
            Remote::Tcp(stream) => stream,
            Remote::Udp(socket, _) => socket,
        }
    }
}

enum Resource {
    Listener(Listener),
    Remote(Remote),
}

impl Resource {
    pub fn event_source(&mut self) -> &mut dyn event::Source {
        match self {
            Resource::Listener(listener) => listener.event_source(),
            Resource::Remote(remote) => remote.event_source(),
        }
    }

    pub fn local_addr(&self) -> SocketAddr {
        match self {
            Resource::Listener(listener) => listener.local_addr(),
            Resource::Remote(remote) => remote.local_addr(),
        }
    }
}

impl std::fmt::Display for Resource {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let resource = match self {
            Resource::Listener(listener) => match listener {
                Listener::Tcp(_) => "Listener::Tcp",
                Listener::Udp(_) => "Listener::Udp",
            },
            Resource::Remote(remote) => match remote {
                Remote::Tcp(_) => "Remote::Tcp",
                Remote::Udp(_, _) => "Remote::Udp",
            },
        };
        write!(f, "{}", resource)
    }
}

pub fn adapter() -> (Arc<Mutex<Controller>>, Receiver) {
    let poll = Poll::new().unwrap();
    let controller = Controller::new(poll.registry().try_clone().unwrap());
    let thread_safe_controller = Arc::new(Mutex::new(controller));
    (thread_safe_controller.clone(), Receiver::new(thread_safe_controller, poll))
}

pub struct Controller {
    resources: HashMap<usize, Resource>,
    last_id: usize,
    registry: Registry,
}

impl Controller {
    fn new(registry: Registry) -> Controller {
        Controller { resources: HashMap::new(), last_id: 0, registry }
    }

    fn add_resource<S: event::Source + ?Sized>(&mut self, source: &mut S) -> usize {
        let id = self.last_id;
        self.last_id += 1;
        self.registry.register(source, Token(id), Interest::READABLE).unwrap();
        id
    }

    pub fn add_remote(&mut self, mut remote: Remote) -> Endpoint {
        let id = self.add_resource(remote.event_source());
        let endpoint = Endpoint::new(id, remote.peer_addr());
        self.resources.insert(id, Resource::Remote(remote));
        endpoint
    }

    pub fn add_listener(&mut self, mut listener: Listener) -> (usize, SocketAddr) {
        let id = self.add_resource(listener.event_source());
        let local_addr = listener.local_addr();
        self.resources.insert(id, Resource::Listener(listener));
        (id, local_addr)
    }

    pub fn remove_resource(&mut self, resource_id: usize) -> Option<()> {
        if let Some(mut resource) = self.resources.remove(&resource_id) {
            self.registry.deregister(resource.event_source()).unwrap();
            Some(())
        }
        else {
            None
        }
    }

    pub fn local_address(&mut self, resource_id: usize) -> Option<SocketAddr> {
        if let Some(resource) = self.resources.get(&resource_id) {
            Some(resource.local_addr())
        }
        else {
            None
        }
    }

    fn send_stream(stream: &mut TcpStream, data: &[u8]) {
        // TODO: The current implementation implies an active waiting,
        // improve it using POLLIN instead to avoid active waiting.
        let mut total_bytes_sent = 0;
        loop {
            match stream.write(&data[total_bytes_sent..]) {
                Ok(bytes_sent) => {
                    total_bytes_sent += bytes_sent;
                    if total_bytes_sent == data.len() {
                        break
                    }
                    // We get sending to data, but not the totality.
                    // We start waiting actively.
                }
                // If WouldBlock is received in this non-blocking socket means that
                // the sending buffer is full and it should wait to send more data.
                // This occurs when huge amounts of data are sent and It could be
                // intensified if the remote endpoint reads slower than this enpoint sends.
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => continue,
                // Skipping. Others errors will be considered fatal for the connection.
                // We skip here their handling because if the connection brokes,
                // an Event::Disconnection will be generated later.
                Err(_) => break, // Err(err).unwrap()
            }
        }
    }

    fn send_datagram(data_len: usize, mut send: impl FnMut() -> io::Result<usize>) {
        if data_len > MAX_UDP_LEN {
            panic!(
                "The datagram max size is {}, your message data takes {}. \
                Split the message in several messages or use an stream protocol as TCP",
                MAX_UDP_LEN,
                data_len
            );
        }

        match send() {
            // The datagram always fit in the MTU
            Ok(_) => (),
            // In this context means that UDP packet that exceeds MTU size.
            // Since this is managing by MAX_UDP_LEN. It should not ocurr.
            // Always can means in that send method is called without knowing
            // the remote addr.
            Err(err) => Err(err).expect("To not occur"),
        }
    }

    pub fn send(&mut self, endpoint: Endpoint, data: &[u8]) -> io::Result<()> {
        if let Some(resource) = self.resources.get_mut(&endpoint.resource_id()) {
            match resource {
                Resource::Listener(listener) => match listener {
                    Listener::Udp(socket) => Self::send_datagram(data.len(), || {
                        socket.send_to(data, endpoint.addr())
                    }),
                    _ => unreachable!(),
                },
                Resource::Remote(remote) => match remote {
                    Remote::Tcp(stream) => Self::send_stream(stream, data),
                    Remote::Udp(socket, _) => Self::send_datagram(data.len(), || {
                        socket.send(data)
                    }),
                },
            }
            Ok(())
        }
        else {
            //TODO: should panics
            Err(io::Error::new(
                ErrorKind::NotFound,
                format!(
                    "Resource id '{}' not exists in the network adapter",
                    endpoint.resource_id()
                ),
            ))
        }
    }
}

pub struct Receiver {
    controller: Arc<Mutex<Controller>>,
    poll: Poll,
    events: Events,
}

impl<'a> Receiver {
    fn new(controller: Arc<Mutex<Controller>>, poll: Poll) -> Receiver {
        Receiver { controller, poll, events: Events::with_capacity(EVENTS_SIZE) }
    }

    pub fn receive<C>(
        &mut self,
        input_buffer: &mut [u8],
        timeout: Option<Duration>,
        event_callback: C,
    ) where
        C: for<'b> FnMut(Endpoint, Event<'b>),
    {
        loop {
            match self.poll.poll(&mut self.events, timeout) {
                Ok(_) => break self.process_event(input_buffer, event_callback),
                Err(e) => match e.kind() {
                    ErrorKind::Interrupted => continue,
                    _ => Err(e).unwrap(),
                },
            }
        }
    }

    fn process_event<C>(&mut self, input_buffer: &mut [u8], mut event_callback: C)
    where C: for<'b> FnMut(Endpoint, Event<'b>) {
        for mio_event in &self.events {
            let token = mio_event.token();
            let id = token.0;
            let mut controller = self.controller.lock().unwrap();

            let resource = controller.resources.get_mut(&id).unwrap();
            log::trace!("Wake from poll for endpoint {}. Resource: {}", id, resource);
            match resource {
                Resource::Listener(listener) => match listener {
                    Listener::Tcp(listener) => {
                        let mut listener = listener;
                        loop {
                            match listener.accept() {
                                Ok((stream, _)) => {
                                    let endpoint = controller.add_remote(Remote::Tcp(stream));
                                    event_callback(endpoint, Event::Connection);

                                    // Used to avoid the consecutive mutable borrows
                                    listener = match controller.resources.get_mut(&id).unwrap() {
                                        Resource::Listener(Listener::Tcp(listener)) => listener,
                                        _ => unreachable!(),
                                    }
                                }
                                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => break,
                                Err(ref err) if err.kind() == io::ErrorKind::Interrupted => {
                                    continue
                                }
                                Err(err) => Err(err).unwrap(),
                            }
                        }
                    }
                    Listener::Udp(socket) => loop {
                        match socket.recv_from(input_buffer) {
                            Ok((size, addr)) => event_callback(
                                Endpoint::new(id, addr),
                                Event::Data(&input_buffer[..size]),
                            ),
                            Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => break,
                            Err(err) => Err(err).unwrap(),
                        }
                    },
                },
                Resource::Remote(remote) => match remote {
                    Remote::Tcp(stream) => loop {
                        match stream.read(input_buffer) {
                            Ok(0) => {
                                let endpoint = Endpoint::new(id, stream.peer_addr().unwrap());
                                controller.remove_resource(endpoint.resource_id()).unwrap();
                                event_callback(endpoint, Event::Disconnection);
                                break
                            }
                            Ok(size) => {
                                let endpoint = Endpoint::new(id, stream.peer_addr().unwrap());
                                event_callback(endpoint, Event::Data(&input_buffer[..size]));
                            }
                            Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => break,
                            Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue,
                            Err(err) => Err(err).unwrap(),
                        }
                    },
                    Remote::Udp(socket, addr) => loop {
                        match socket.recv(input_buffer) {
                            Ok(size) => event_callback(
                                Endpoint::new(id, *addr),
                                Event::Data(&input_buffer[..size]),
                            ),
                            Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => break,
                            Err(ref err) if err.kind() == io::ErrorKind::ConnectionRefused => {
                                continue
                            }
                            Err(err) => Err(err).unwrap(),
                        }
                    },
                },
            }
        }
    }
}