#![cfg_attr(target_os = "yggdrasil", feature(yggdrasil_os, rustc_private))]
#![feature(if_let_guard)]
use std::{
    collections::HashSet,
    io::{self, Read, Write},
    net::SocketAddr,
    os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd},
    path::PathBuf,
    process::{self, Command, ExitCode, Stdio},
    str::FromStr,
    time::Duration,
};

use clap::Parser;
use cross::{
    io::{PidFd, PtyMaster, TerminalOptions, TerminalOptionsImpl, TerminalSize},
    process::CommandSpawnExt,
};
use rsh::{
    crypt::config::{ServerConfig, SimpleServerKeyStore},
    proto::TerminalInfo,
    server::{self, Server, Session, SessionClient},
};

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

#[derive(Debug, clap::Parser)]
struct Args {
    #[clap(short = 'P', long, help = "rsh listen port", default_value_t = 77)]
    port: u16,
    #[clap(
        short = 'S',
        long,
        help = "where rsh will load private keys from",
        default_value = "/etc/rsh"
    )]
    keystore: PathBuf,
}

pub struct SessionImpl {
    fds: [RawFd; 2],
    pty_master: PtyMaster,
    remote: SocketAddr,
    shell: process::Child,
    pidfd: Option<PidFd>,
}

impl Session for SessionImpl {
    type Error = std::io::Error;

    fn open(remote: &SocketAddr, terminal: &TerminalInfo) -> Result<Self, Self::Error> {
        let remote = *remote;

        let termios = TerminalOptionsImpl::normal();
        let (pty_master, pty_slave) = cross::io::open_pty(
            &termios,
            TerminalSize {
                rows: terminal.rows as _,
                columns: terminal.columns as _,
                x_pixels: 0,
                y_pixels: 0,
            },
        )?;

        let pty_slave_stdin = pty_slave.into_raw_fd();
        let pty_slave_stdout = cross::io::clone_fd(pty_slave_stdin)?;
        let pty_slave_stderr = cross::io::clone_fd(pty_slave_stdin)?;

        let shell = unsafe {
            Command::new("/bin/sh")
                .arg("-l")
                .stdin(Stdio::from_raw_fd(pty_slave_stdin))
                .stdout(Stdio::from_raw_fd(pty_slave_stdout))
                .stderr(Stdio::from_raw_fd(pty_slave_stderr))
                .create_session()?
                .spawn()?
        };
        let pidfd = PidFd::new(shell.id())?;

        let fds = [pty_master.as_raw_fd(), pidfd.as_raw_fd()];

        Ok(Self {
            remote,
            pty_master,
            pidfd: Some(pidfd),
            shell,
            fds,
        })
    }

    fn close(mut self) -> Result<(), Self::Error> {
        self.shell.wait()?;
        self.pidfd = None;
        Ok(())
    }

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

    fn handle_input(&mut self, input: &[u8], _client: SessionClient) -> Result<bool, Self::Error> {
        self.pty_master.write_all(input)?;
        Ok(false)
    }

    fn read_output(
        &mut self,
        fd: std::os::fd::RawFd,
        buffer: &mut [u8],
    ) -> Result<Option<usize>, Self::Error> {
        if fd == self.fds[0] {
            self.pty_master.read(buffer).map(Some)
        } else if fd == self.fds[1] {
            let Some(pidfd) = self.pidfd.as_mut() else {
                return Ok(None);
            };

            match pidfd.exit_status() {
                Ok(status) => {
                    log::info!("Shell exited with status: {status}");
                    Ok(Some(0))
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => Ok(None),
                Err(e) => Err(e),
            }
        } else {
            unreachable!()
        }
    }

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

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Server error: {0}")]
    Server(#[from] server::Error),
}

fn run(args: Args) -> Result<(), Error> {
    let keystore = SimpleServerKeyStore {
        path: args.keystore,
        accepted_keys: HashSet::from_iter([
            "ed25519 sha256 ab:5d:d8:a9:56:79:f9:0f:47:32:a4:4b:f5:fb:49:2d:3d:91:a7:7d:70:60:8a:57:84:36:43:80:7e:91:8a:c3".into()
        ]),
    };
    let server_config = ServerConfig::with_default_algorithms(keystore);
    let listen_addr = SocketAddr::from_str(&format!("0.0.0.0:{}", args.port)).unwrap();
    let server = Server::<SessionImpl>::listen(listen_addr, server_config)?;
    server.run().map_err(From::from)
}

fn main() -> ExitCode {
    logsink::setup_logging(true);
    let args = Args::parse();
    if let Err(error) = run(args) {
        eprintln!("Finished with error: {error}");
        ExitCode::FAILURE
    } else {
        ExitCode::SUCCESS
    }
}
