use std::{
    collections::{hash_map::Entry, HashMap},
    fmt,
    hint::unreachable_unchecked,
    io::{self, Read, Write},
    net::SocketAddr,
    os::fd::{AsRawFd, RawFd},
    process::{Child, Command},
    time::Duration,
};

use cross::io::{PidFd, Pipe, Poll};

use crate::{
    crypt::{
        client,
        config::ServerConfig,
        server::{self, ServerSocket},
        stream::{self, ClientSocket},
    },
    proto::{ClientMessage, ServerMessage, StreamIndex, TerminalInfo},
};

pub const PING_INTERVAL: Duration = Duration::from_millis(500);

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),
    #[error("Socket error: {0}")]
    Socket(#[from] server::Error),
    #[error("Client error: {0}")]
    Client(#[from] client::Error),
    #[error("Client stream error: {0}")]
    Stream(#[from] stream::Error),
}

pub trait Session: Sized {
    type Error: fmt::Display;

    fn open(peer: &SocketAddr, terminal: &TerminalInfo) -> Result<Self, Self::Error>;
    fn peer(&self) -> SocketAddr;
    fn handle_input(&mut self, input: &[u8], client: SessionClient) -> Result<bool, Self::Error>;
    fn read_output(&mut self, fd: RawFd, buffer: &mut [u8]) -> Result<Option<usize>, Self::Error>;
    fn event_fds(&self) -> &[RawFd];
    fn close(self) -> Result<(), Self::Error>;
}

pub struct SessionClient<'s> {
    transport: &'s mut ClientSocket,
}

pub struct EchoSession {
    peer: SocketAddr,
}

struct PendingCommand {
    child: Child,
    child_pid: Option<PidFd>,
    stdin: Option<Pipe>,
    stdout: Option<Pipe>,
    stderr: Option<Pipe>,
}

enum ClientSession<T: Session> {
    None,
    Terminal(T),
    Command(PendingCommand),
}

struct Client<T: Session> {
    stream: ClientSocket,
    session: ClientSession<T>,
}

struct ClientSet<T: Session> {
    last_session_key: u64,
    clients: HashMap<u64, Client<T>>,
    socket_fd_map: HashMap<RawFd, u64>,
    session_fd_map: HashMap<RawFd, u64>,
}

pub struct Server<T: Session> {
    poll: Poll,
    socket: ServerSocket,

    clients: ClientSet<T>,
}

enum CommandEvent {
    Output(StreamIndex, Result<usize, io::Error>),
    Exited(Result<i32, io::Error>),
}

impl<T: Session> Client<T> {
    fn make_command(text: &str) -> Result<PendingCommand, Error> {
        let mut words = text.split(' ');
        let program = words.next().unwrap();
        let mut command = Command::new(program);

        let (stdin_read, stdin_write) = Pipe::new(false, false)?;
        let (stdout_read, stdout_write) = Pipe::new(true, false)?;
        let (stderr_read, stderr_write) = Pipe::new(true, false)?;

        command
            .args(words)
            .stdin(stdin_read.to_child_stdio())
            .stdout(stdout_write.to_child_stdio())
            .stderr(stderr_write.to_child_stdio());

        let child = command.spawn()?;
        let child_pid = PidFd::new(child.id())?;

        Ok(PendingCommand {
            child,
            child_pid: Some(child_pid),
            stdout: Some(stdout_read),
            stderr: Some(stderr_read),
            stdin: Some(stdin_write),
        })
    }
}

impl<T: Session> ClientSet<T> {
    pub fn new() -> Self {
        Self {
            last_session_key: 1,
            clients: HashMap::new(),
            socket_fd_map: HashMap::new(),
            session_fd_map: HashMap::new(),
        }
    }

    pub fn add_client(&mut self, stream: ClientSocket, poll: &mut Poll) -> Result<(), Error> {
        let (key, client) = loop {
            let key = self.last_session_key;
            self.last_session_key += 1;

            match self.clients.entry(key) {
                Entry::Vacant(entry) => {
                    let session = entry.insert(Client {
                        stream,
                        session: ClientSession::None,
                    });
                    break (key, session);
                }
                Entry::Occupied(_) => continue,
            }
        };

        log::info!("Add client {key}");
        poll.add(&client.stream)?;
        self.socket_fd_map.insert(client.stream.as_raw_fd(), key);

        Ok(())
    }

    fn remove(&mut self, key: u64, poll: &mut Poll) -> Result<(), Error> {
        if let Some(client) = self.clients.remove(&key) {
            poll.remove(&client.stream)?;
            self.socket_fd_map.remove(&client.stream.as_raw_fd());

            match client.session {
                ClientSession::Terminal(terminal) => {
                    for fd in terminal.event_fds() {
                        log::info!("{key}: remove fd {fd:?}");
                        poll.remove(fd)?;
                        self.session_fd_map.remove(fd);
                    }
                }
                ClientSession::Command(mut command) => {
                    if let Some(stdout) = command.stdout.take() {
                        poll.remove(&stdout)?;
                        self.session_fd_map.remove(&stdout.as_raw_fd());
                    }
                    if let Some(stderr) = command.stderr.take() {
                        poll.remove(&stderr)?;
                        self.session_fd_map.remove(&stderr.as_raw_fd());
                    }
                    if let Some(child_pid) = command.child_pid.take() {
                        poll.remove(&child_pid)?;
                        self.session_fd_map.remove(&child_pid.as_raw_fd());
                    }
                }
                ClientSession::None => (),
            }

            log::info!("Remove client {key}");
        }
        Ok(())
    }

    pub fn handle_input(&mut self, fd: RawFd, poll: &mut Poll) -> Result<(), Error> {
        let mut buffer = [0; 512];

        if let Some(&key) = self.socket_fd_map.get(&fd) {
            let client = self.clients.get_mut(&key).unwrap();
            let peer = client.stream.remote_address();

            let mut closed = match client.stream.poll() {
                Ok(0) => true,
                Ok(_) => false,
                Err(error) => {
                    log::error!("{peer}: {error}");
                    return self.remove(key, poll);
                }
            };

            loop {
                let message = match client.stream.read(&mut buffer) {
                    Ok(Some(message)) => message,
                    Ok(None) => break,
                    Err(error) => {
                        log::error!("{peer}: {error}");
                        closed = true;
                        break;
                    }
                };

                match (&message, &mut client.session) {
                    (ClientMessage::RunCommand(command), ClientSession::None) => {
                        log::info!("{peer}: cmd {command:?}");

                        let command = match Client::<T>::make_command(command) {
                            Ok(command) => command,
                            Err(error) => {
                                log::error!("{peer}: command error: {error}");
                                return self.remove(key, poll);
                            }
                        };

                        let stdout_fd = command.stdout.as_ref().unwrap().as_raw_fd();
                        let stderr_fd = command.stderr.as_ref().unwrap().as_raw_fd();
                        let child_pid_fd = command.child_pid.as_ref().unwrap().as_raw_fd();

                        poll.add(&stdout_fd)?;
                        poll.add(&stderr_fd)?;
                        poll.add(&child_pid_fd)?;

                        self.session_fd_map.insert(stdout_fd, key);
                        self.session_fd_map.insert(stderr_fd, key);
                        self.session_fd_map.insert(child_pid_fd, key);

                        client.session = ClientSession::Command(command);
                    }
                    (ClientMessage::OpenSession(terminal), ClientSession::None) => {
                        log::info!("{peer}: new session");
                        let terminal = match T::open(&peer, terminal) {
                            Ok(session) => session,
                            Err(error) => {
                                log::error!("{peer}: couldn't open session: {error}");
                                self.remove(key, poll)?;
                                return Ok(());
                            }
                        };
                        let terminal = client.session.set_terminal(terminal);
                        for fd in terminal.event_fds() {
                            log::info!("{key}: add fd {fd:?}");
                            poll.add(fd)?;
                            self.session_fd_map.insert(*fd, key);
                        }

                        client.stream.write_all(&ServerMessage::SessionOpen).ok();
                    }
                    (ClientMessage::Input(data), ClientSession::Command(command)) => {
                        if let Some(stdin) = command.stdin.as_mut() {
                            match stdin.write_all(data) {
                                Ok(()) => {
                                    stdin.flush().ok();
                                }
                                Err(error) => {
                                    log::error!("{peer}: stdin error: {error}");
                                    return self.remove(key, poll);
                                }
                            }
                        } else {
                            log::warn!("{peer}: stdin already closed");
                        }
                    }
                    (ClientMessage::Input(data), ClientSession::Terminal(terminal)) => {
                        let client = SessionClient {
                            transport: &mut client.stream,
                        };
                        match terminal.handle_input(data, client) {
                            Ok(false) => (),
                            Ok(true) => {
                                log::info!("{peer}: session closed");
                                return self.remove(key, poll);
                            }
                            Err(error) => {
                                log::info!("{peer}: session error: {error}");
                                return self.remove(key, poll);
                            }
                        }
                    }
                    (ClientMessage::CloseStdin, ClientSession::Command(command)) => {
                        // Drop stdin handle
                        command.stdin.take();
                    }
                    _ => {
                        log::warn!("{peer}: unhandled message");
                    }
                }
            }

            if closed {
                log::info!("{peer}: stream closed");
                self.remove(key, poll)?;
            }

            Ok(())
        } else if let Some(&key) = self.session_fd_map.get(&fd) {
            let client = self.clients.get_mut(&key).unwrap();
            let peer = client.stream.remote_address();

            let (mut len, stream_index) = match &mut client.session {
                ClientSession::Command(command) => match command.read_output(fd, &mut buffer) {
                    CommandEvent::Output(index, Ok(len)) => {
                        if len == 0 {
                            poll.remove(&fd)?;
                            self.socket_fd_map.remove(&fd);

                            match index {
                                StreamIndex::Stdout => {
                                    command.stdout = None;
                                }
                                StreamIndex::Stderr => {
                                    command.stderr = None;
                                }
                            }

                            if command.is_dead() {
                                log::info!("{peer}: command stdout/stderr closed");
                                return self.remove(key, poll);
                            }
                        }
                        (len, index)
                    }
                    CommandEvent::Output(index, Err(error)) => {
                        if error.kind() == io::ErrorKind::WouldBlock {
                            return Ok(());
                        }
                        log::error!("{peer}: {index:?} error: {error}");
                        return self.remove(key, poll);
                    }
                    CommandEvent::Exited(status) => {
                        log::info!("{peer}: command exited: {status:?}");
                        poll.remove(&fd)?;
                        self.socket_fd_map.remove(&fd);
                        command.child_pid = None;

                        if command.is_dead() {
                            log::info!("{peer}: command stdout/stderr closed");
                            return self.remove(key, poll);
                        }

                        return Ok(());
                    }
                },
                ClientSession::Terminal(terminal) => match terminal.read_output(fd, &mut buffer) {
                    Ok(None) => {
                        return Ok(());
                    }
                    Ok(Some(0)) => {
                        poll.remove(&fd)?;
                        self.socket_fd_map.remove(&fd);

                        // TODO check for process as well
                        log::info!("{peer}: terminal closed");
                        return self.remove(key, poll);
                    }
                    Ok(Some(len)) => (len, StreamIndex::Stdout),
                    Err(error) => {
                        log::error!("{peer}: terminal error: {error}");
                        return self.remove(key, poll);
                    }
                },
                ClientSession::None => unreachable!(),
            };

            if len == 0 {
                log::info!("{peer}: {stream_index:?} closed");
                return Ok(());
                // return self.remove(key, poll);
            }

            // Split output into 128-byte chunks
            let mut pos = 0;
            while len != 0 {
                let amount = core::cmp::min(len, 128);

                if let Err(error) = client.stream.write_all(&ServerMessage::Output(
                    stream_index,
                    &buffer[pos..pos + amount],
                )) {
                    log::error!("{peer}: communication error: {error}");
                    return self.remove(key, poll);
                }

                pos += amount;
                len -= amount;
            }

            Ok(())
        } else {
            unreachable!()
        }
    }
}

impl PendingCommand {
    pub fn read_output(&mut self, fd: RawFd, buffer: &mut [u8]) -> CommandEvent {
        if let Some(stdout) = self.stdout.as_mut()
            && fd == stdout.as_raw_fd()
        {
            return CommandEvent::Output(StreamIndex::Stdout, stdout.read(buffer));
        }
        if let Some(stderr) = self.stderr.as_mut()
            && fd == stderr.as_raw_fd()
        {
            return CommandEvent::Output(StreamIndex::Stderr, stderr.read(buffer));
        }
        if let Some(child_pid) = self.child_pid.as_mut()
            && fd == child_pid.as_raw_fd()
        {
            let status = child_pid.exit_status();
            return CommandEvent::Exited(status);
        }
        unreachable!()
    }

    pub fn is_dead(&self) -> bool {
        self.child_pid.is_none() && self.stdout.is_none() && self.stderr.is_none()
    }
}

impl Drop for PendingCommand {
    fn drop(&mut self) {
        self.child.wait().ok();
    }
}

impl<T: Session> ClientSession<T> {
    pub fn set_terminal(&mut self, terminal: T) -> &mut T {
        *self = Self::Terminal(terminal);
        match self {
            Self::Terminal(terminal) => terminal,
            _ => unsafe { unreachable_unchecked() },
        }
    }
}

impl<T: Session> Server<T> {
    pub fn listen(listen_addr: SocketAddr, crypto_config: ServerConfig) -> Result<Self, Error> {
        let mut poll = Poll::new()?;
        let socket = ServerSocket::bind(listen_addr, crypto_config)?;
        poll.add(&socket)?;
        log::info!("Listening on {listen_addr}");
        Ok(Self {
            poll,
            socket,
            clients: ClientSet::new(),
        })
    }

    pub fn run(mut self) -> Result<(), Error> {
        loop {
            let fd = self.poll.wait(None)?.unwrap();

            match fd {
                _ if fd == self.socket.as_raw_fd() => {
                    // Poll server socket
                    match self.socket.poll()? {
                        Some(client) => {
                            self.clients.add_client(client, &mut self.poll)?;
                        }
                        None => continue,
                    }
                }
                _ => {
                    // Client/Session-related activity
                    self.clients.handle_input(fd, &mut self.poll)?;
                }
            }
        }
    }
}

impl SessionClient<'_> {
    pub fn send_data(&mut self, data: &[u8]) -> Result<(), Error> {
        self.send_message(&ServerMessage::Output(StreamIndex::Stdout, data))
    }

    pub fn send_message(&mut self, message: &ServerMessage) -> Result<(), Error> {
        self.transport.write_all(message).map_err(Error::from)
    }
}

impl Session for EchoSession {
    type Error = Error;

    fn open(peer: &SocketAddr, _terminal: &TerminalInfo) -> Result<Self, Self::Error> {
        Ok(Self { peer: *peer })
    }

    fn close(self) -> Result<(), Self::Error> {
        Ok(())
    }

    fn event_fds(&self) -> &[RawFd] {
        &[]
    }

    fn read_output(
        &mut self,
        _fd: RawFd,
        _buffer: &mut [u8],
    ) -> Result<Option<usize>, Self::Error> {
        Ok(None)
    }

    fn handle_input(
        &mut self,
        input: &[u8],
        mut client: SessionClient,
    ) -> Result<bool, Self::Error> {
        if input.contains(&b'\x04') {
            return Ok(true);
        }
        log::debug!("{:02x?}", input);
        client.send_data(input)?;
        Ok(false)
    }

    fn peer(&self) -> SocketAddr {
        self.peer
    }
}
