use std::{
    ffi::OsStr,
    fs::File,
    io::{self, Read, Write},
    net::ToSocketAddrs,
    path::{Path, PathBuf},
    process::ExitCode,
    str::FromStr,
};

use clap::Parser;
use cross::fs::{MetadataExt, TempFile};
use rsh::{
    crypt::{
        self, client,
        config::{ClientConfig, SimpleClientKeyStore},
        signature::{SignEd25519, SignatureMethod},
        stream::{ClientSocket, Message},
    },
    proto::{
        ClientMessage, ClientTransferControl, RequestOpenWrite, ServerMessage,
        ServerTransferControl,
    },
};

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("rsh client error: {0}")]
    Client(#[from] client::Error),
    #[error("rsh stream error: {0}")]
    Stream(#[from] crypt::stream::Error),
    #[error("rsh signature error: {0}")]
    Signature(#[from] crypt::signature::Error),
    #[error("unresolved hostname: {0}")]
    UnresolvedHostname(String),
    #[error("invalid path: {0:?}")]
    InvalidPath(PathBuf),
    #[error("remote closed connection: {0:?}")]
    RemoteClosedConnection(String),
    #[error("transfer error: {0:?}")]
    Transfer(String),
}

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, long)]
    key: PathBuf,
    #[clap(short = 'P', long, default_value_t = 77)]
    port: u16,

    src: ResourcePath,
    dst: ResourcePath,
}

#[derive(Clone, Debug)]
enum ResourcePath {
    Local(PathBuf),
    Remote(String, PathBuf),
}

enum Outcome {
    Write(Vec<u8>),
    Read(u32),
    CloseDisconnect,
    None,
}

fn run_loop<F>(
    remote_host: &str,
    remote_port: u16,
    config: ClientConfig,
    init: ClientTransferControl,
    mut handler: F,
) -> Result<(), Error>
where
    F: FnMut(&ServerTransferControl) -> Result<Outcome, Error>,
{
    let mut remote_hosts = (remote_host, remote_port).to_socket_addrs()?;
    let remote_host = remote_hosts
        .next()
        .ok_or(Error::UnresolvedHostname(remote_host.to_owned()))?;
    let mut socket = ClientSocket::connect(remote_host, config)?;
    let mut buffer = [0; 512];

    socket.write_all(&ClientMessage::TransferControl(init))?;

    loop {
        let message: ServerMessage = match socket.poll_read(&mut buffer)? {
            Message::Data(data) => data,
            Message::Incomplete => continue,
            Message::Closed => break,
        };

        if let ServerMessage::Bye(reason) = &message {
            return Err(Error::RemoteClosedConnection((*reason).to_owned()));
        }

        let ServerMessage::TransferControl(control) = message else {
            log::warn!("Unexpected message: {message:?}");
            continue;
        };

        if let ServerTransferControl::TransferError(error) = control {
            socket.write_all(&ClientMessage::Bye("disconnect"))?;
            return Err(Error::Transfer(error.into()));
        }

        match handler(&control)? {
            Outcome::CloseDisconnect => {
                socket.write_all(&ClientMessage::TransferControl(
                    ClientTransferControl::RequestClose,
                ))?;
                socket.write_all(&ClientMessage::Bye("disconnect"))?;
                break;
            }
            Outcome::Read(amount) => {
                socket.write_all(&ClientMessage::TransferControl(
                    ClientTransferControl::RequestRead(amount),
                ))?;
            }
            Outcome::Write(data) => {
                socket.write_all(&ClientMessage::TransferControl(
                    ClientTransferControl::RequestWrite(&data),
                ))?;
            }
            Outcome::None => (),
        }
    }

    Ok(())
}

impl FromStr for ResourcePath {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Some((host, filename)) = s.split_once(':') {
            Ok(ResourcePath::Remote(host.into(), filename.into()))
        } else {
            Ok(ResourcePath::Local(s.into()))
        }
    }
}

fn make_config<P: AsRef<Path>>(key: P) -> Result<ClientConfig, Error> {
    let ed25519 = SignEd25519::load_signing_key(key)?;
    let key = SignatureMethod::Ed25519(ed25519);
    let config = ClientConfig::with_default_algorithms(SimpleClientKeyStore::new(key));
    Ok(config)
}

fn copy_local_to_remote(
    args: &Args,
    config: ClientConfig,
    local_path: &Path,
    remote_host: &str,
    remote_path: &Path,
) -> Result<(), Error> {
    let Some(local_filename) = local_path.file_name().and_then(OsStr::to_str) else {
        return Err(Error::InvalidPath(local_path.to_owned()));
    };
    let remote_path = remote_path.to_str().unwrap();
    let file = File::open(local_path)?;
    let metadata = file.metadata()?;
    let mut file = Some(file);

    log::info!("copy local to remote: {local_path:?} -> {remote_host}, {remote_path:?}");

    let mode = metadata.mode();
    let open_info = RequestOpenWrite {
        filename: local_filename,
        path: remote_path,
        mode,
        create: true,
    };
    let mut buffer = [0; 256];

    run_loop(
        remote_host,
        args.port,
        config,
        ClientTransferControl::RequestOpenWrite(open_info),
        |packet| match (packet, file.as_mut()) {
            (ServerTransferControl::WriteOpened, Some(f)) => {
                let len = f.read(&mut buffer)?;
                if len == 0 {
                    file = None;
                }
                Ok(Outcome::Write(buffer[..len].into()))
            }
            (&ServerTransferControl::WriteAck(0), None) => Ok(Outcome::CloseDisconnect),
            (&ServerTransferControl::WriteAck(_count), Some(f)) => {
                // TODO handle ACKed count
                let len = f.read(&mut buffer)?;
                if len == 0 {
                    file = None;
                }
                Ok(Outcome::Write(buffer[..len].into()))
            }
            _ => Ok(Outcome::None),
        },
    )?;

    Ok(())
}

fn copy_remote_to_local(
    args: &Args,
    config: ClientConfig,
    local_path: &Path,
    remote_host: &str,
    remote_path: &Path,
) -> Result<(), Error> {
    let Some(remote_filename) = remote_path.file_name().and_then(OsStr::to_str) else {
        return Err(Error::InvalidPath(remote_path.to_owned()));
    };
    let remote_path = remote_path.to_str().unwrap();
    let local_path = if local_path.is_dir() {
        local_path.join(remote_filename)
    } else {
        local_path.to_owned()
    };
    let local_temp_path = local_path.with_added_extension("tmp");
    let mut temp = TempFile::new(&local_temp_path)?;
    let mut read_position = None;

    log::info!("copy remote to local: {remote_host}, {remote_path:?} -> {local_path:?}");

    run_loop(
        remote_host,
        args.port,
        config,
        ClientTransferControl::RequestOpenRead(remote_path),
        |packet| match (packet, read_position) {
            (ServerTransferControl::ReadOpened(_info), None) => {
                read_position = Some(0);
                Ok(Outcome::Read(384))
            }
            (ServerTransferControl::ReadData(data), Some(pos)) => {
                if data.is_empty() {
                    // EOF
                    Ok(Outcome::CloseDisconnect)
                } else {
                    temp.write_all(data)?;
                    read_position = Some(pos + data.len());
                    Ok(Outcome::Read(384))
                }
            }
            _ => Ok(Outcome::None),
        },
    )?;

    temp.rename(&local_path)?;

    Ok(())
}

fn run(args: Args) -> Result<(), Error> {
    let config = make_config(&args.key)?;
    match (&args.src, &args.dst) {
        (ResourcePath::Local(local), ResourcePath::Remote(host, remote)) => {
            copy_local_to_remote(&args, config, local.as_ref(), host, remote.as_ref())
        }
        (ResourcePath::Remote(host, remote), ResourcePath::Local(local)) => {
            copy_remote_to_local(&args, config, local.as_ref(), host, remote.as_ref())
        }
        _ => todo!(),
    }
}

fn main() -> ExitCode {
    logsink::setup_logging(true);
    let args = Args::parse();
    match run(args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("rcp: {error}");
            ExitCode::FAILURE
        }
    }
}
