use std::collections::HashMap;
use std::fs::{OpenOptions, File};
use std::io::Write;
use std::pin::Pin;

use futures::FutureExt;
use log::info;
use rand::{Rng, thread_rng};
use russh::server::{Auth, Session};
use russh::*;

fn send_str(session: &mut Session, channel: ChannelId, s: String) {
    let data = CryptoVec::from(s);
    session.data(channel, data);
}

enum ChannelState {
    Shell,
    Scp,
    Forwarding,
}

enum CommandStatus {
    Done,
    Scp,
}

const UNAME_RESPONSES: &[&str] = &[
    "Linux fnordister 5.10.0-10-amd64 #1 SMP Debian 5.10.84-1 (2021-12-08) x86_64 GNU/Linux",
    "Linux raspberrypi 5.19.9 #1-NixOS SMP Thu Sep 15 08:47:20 UTC 2022 aarch64 GNU/Linux",
    "Linux OpenWrt 5.10.138 #0 Sat Sep 3 02:55:34 2022 mips GNU/Linux",
];

pub struct Handler {
    filename: String,
    lazy_file: Option<File>,
    buffer: Vec<u8>,
    user: String,
    channels: HashMap<ChannelId, ChannelState>,
}

impl Handler {
    pub fn new<P: ToString>(path: P) -> Self {
        Handler {
            filename: path.to_string(),
            lazy_file: None,
            buffer: vec![],
            user: "root".into(),
            channels: HashMap::new(),
        }
    }

    fn file(&mut self) -> &mut File {
        if self.lazy_file.is_none() {
            let file = OpenOptions::new()
                .create(true)
                .append(true)
                .open(&self.filename)
                .unwrap();
            self.lazy_file = Some(file);
        }
        self.lazy_file.as_mut().unwrap()
    }

    fn send_prompt(&self, session: &mut Session, channel: ChannelId) {
        send_str(session, channel, format!("{}@fnordister:~$ ", self.user));
    }

    fn handle_command<F: FnMut(&str)>(&self, command: String, mut respond: F) -> CommandStatus {
        let program_len = command.find(|c: char| c.is_whitespace())
            .unwrap_or(command.len());
        let program = &command[..program_len];

        match program {
            "whoami" => respond(&self.user),
            "id" => respond(&format!("uid=0({}) gid=0(root)", self.user)),
            "uname" => {
                let mut rnd = thread_rng();
                let i = rnd.gen_range(0..UNAME_RESPONSES.len());
                respond(&UNAME_RESPONSES[i]);
            }
            "pwd" => respond("/"),
            "ls" => {
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 .");
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 ..");
                respond("drwxr-xr-x  18 root root    18 Jan  4  0000 ...");
            },
            "scp" => return CommandStatus::Scp,
            "bash" => {}
            "sh" => {}
            "cd" => {}
            "" => {}
            _ => respond(&format!("{}: command not found", program)),
        }

        CommandStatus::Done
    }
}

impl server::Handler for Handler {
    type Error = anyhow::Error;
    type FutureAuth =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Auth)>> + Send>>;
    type FutureUnit =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session)>> + Send>>;
    type FutureBool =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session, bool)>> + Send>>;

    fn finished_auth(self, auth: Auth) -> Self::FutureAuth {
        async { Ok((self, auth)) }.boxed()
    }

    fn finished_bool(self, b: bool, s: Session) -> Self::FutureBool {
        async move { Ok((self, s, b)) }.boxed()
    }

    fn finished(self, s: Session) -> Self::FutureUnit {
        async { Ok((self, s)) }.boxed()
    }

    fn shell_request(mut self, channel: ChannelId, mut session: Session) -> Self::FutureUnit {
        info!("shell_request {}", channel);
        self.channels.insert(channel, ChannelState::Shell);
        self.send_prompt(&mut session, channel);

        self.finished(session)
    }

    fn exec_request(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        info!("exec_request {} {}", channel, String::from_utf8_lossy(data));
        writeln!(self.file(), "Execute: {}\n", String::from_utf8_lossy(data))
            .unwrap();

        let line = String::from_utf8_lossy(data).into();
        let status = self.handle_command(line, |response| {
            let mut data = Vec::from(response);
            data.extend_from_slice(b"\r\n");
            session.data(channel, data.into());
        });
        match status {
            CommandStatus::Scp => {
                session.data(channel, vec![0].into());
                self.channels.insert(channel, ChannelState::Scp);
            }
            CommandStatus::Done => {
                session.close(channel);
            }
        }

        self.finished(session)
    }

    fn subsystem_request(
        mut self,
        channel: ChannelId,
        name: &str,
        session: Session
    ) -> Self::FutureUnit {
        info!("subsystem_request {}", channel);
        writeln!(self.file(), "Subsystem requested: {}\n", name)
            .unwrap();

        self.finished(session)
    }


    /// `ssh -R`
    fn tcpip_forward(mut self, address: &str, port: u32, session: Session) -> Self::FutureBool {
        info!("Reverse-forwarding {}:{}", address, port);
        writeln!(self.file(), "Reverse-forwarding {}:{}", address, port)
            .unwrap();

        let handle = session.handle();
        let address = address.to_string();
        tokio::spawn(async move {
            if let Ok(mut channel) = handle
                .channel_open_forwarded_tcpip(address, port, "0.0.0.0", 1)
                .await {
                    let _ = channel.eof().await;
                }
        });
        self.finished_bool(true, session)
    }

    /// `ssh -L` but doesn't work
    fn channel_open_forwarded_tcpip(
        mut self,
        channel: ChannelId,
        host_to_connect: &str,
        port_to_connect: u32,
        originator_address: &str,
        originator_port: u32,
        session: Session
    ) -> Self::FutureBool {
        info!(
            "Forwarding {}:{} from {}:{}",
            host_to_connect, port_to_connect,
            originator_address, originator_port,
        );
        writeln!(
            self.file(), "Forwarding {}:{} from {}:{}",
            host_to_connect, port_to_connect,
            originator_address, originator_port,
        ).unwrap();

        self.channels.insert(channel, ChannelState::Forwarding);
        self.finished_bool(true, session)
    }

    // fn channel_open_session(
    //     self,
    //     channel: ChannelId,
    //     mut session: Session
    // ) -> Self::FutureBool {
    //     info!("channel_open_session");
    //     // session.channel_success(channel);
    //     self.finished_bool(true, session)
    // }

    // fn channel_open_confirmation(
    //     self,
    //     id: ChannelId,
    //     max_packet_size: u32,
    //     window_size: u32,
    //     mut session: Session
    // ) -> Self::FutureUnit {
    //     self.finished(session)
    // }

    // fn channel_close(self, channel: ChannelId, session: Session) -> Self::FutureUnit {
    //     info!("channel_close {}", channel);
    //     self.finished(session)
    // }

    fn channel_eof(mut self, channel: ChannelId, mut session: Session) -> Self::FutureUnit {
        session.close(channel);
        self.channels.remove(&channel);

        self.finished(session)
    }

    fn auth_password(mut self, user: &str, password: &str) -> Self::FutureAuth {
        writeln!(self.file(), "Authenticated as {} with {}\n", user, password)
            .unwrap();
        self.user = user.into();
        self.finished_auth(server::Auth::Accept)
    }

    fn data(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        self.file().write(data)
            .unwrap();

        match self.channels.get(&channel) {
            Some(ChannelState::Shell) => {
                // echo input back
                session.data(channel, data.to_vec().into());

                if self.buffer.len() < 1024 {
                    self.buffer.extend_from_slice(data);
                }
                if let Some(newline) = self.buffer.iter().position(|b| *b == b'\r') {
                    let rest = self.buffer.split_off(newline + 1);
                    let line: String = String::from_utf8_lossy(&self.buffer).into();
                    self.buffer = rest;

                    session.data(channel, b"\n".to_vec().into());
                    self.handle_command(line, |response| {
                        let mut data = Vec::from(response);
                        data.extend_from_slice(b"\r\n");
                        session.data(channel, data.into());
                    });
                    self.send_prompt(&mut session, channel);
                }
            }

            Some(ChannelState::Scp) => {
                if self.buffer.len() < 1024 {
                    self.buffer.extend_from_slice(data);
                }
                if let Some(newline) = self.buffer.iter().position(|b| *b == b'\n') {
                    let rest = self.buffer.split_off(newline + 1);
                    // let line: String = String::from_utf8_lossy(&self.buffer).into();
                    self.buffer = rest;

                    session.data(channel, vec![0].into());
                }
            }

            Some(ChannelState::Forwarding) => {}

            None => {
                info!("data on unidentified channel {}", channel);
            }
        }

        self.finished(session)
    }

    // fn extended_data(
    //     self,
    //     channel: ChannelId,
    //     code: u32,
    //     data: &[u8],
    //     session: Session
    // ) -> Self::FutureUnit {
    //     info!("extended_data {} {} {:?}", channel, code, data);
    //     self.finished(session)
    // }
}