pub(crate) mod args;

use std::{fs::File, io::Read, ops::Add, path::PathBuf, sync::Arc, time::Duration};

use crate::{commands::args::append_port_if_needed, ssh};
use anyhow::{Context as _, Result as AnyhowResult};
use clap::{Arg, ArgAction, Command};
use deltachat::{
    chat::{self, ChatId},
    contact::ContactId,
    context::Context,
    message::Message,
};
use eui48::MacAddress;
use russh::{client::AuthResult, keys::PrivateKeyWithHashAlg};
use tokio::{
    sync::Mutex,
    time::{Instant, timeout, timeout_at},
};

use crate::{AUTH_REQUIRED, BotContext, config::BotConfig, data_path};

pub async fn echo(dchat_ctx: Arc<Mutex<Context>>, msg: Message) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let dchat_ctx_lock = dchat_ctx.lock().await;
    chat::send_text_msg(&dchat_ctx_lock, chat_id, msg.get_text()).await?;
    Ok(())
}

pub async fn global_usage(ctx: Arc<Mutex<BotContext>>) -> String {
    let ctx_lock = ctx.lock().await;
    let config = &ctx_lock.config;
    let mut usage = String::from("DeltaChat Remote Control Bot\n\nCommands:\n");
    usage += "/auth <password> - Authenticate (prove you are the bot owner or a trusted person)\n";
    usage += &wol_usage(config);
    usage += " - send Wake-on-LAN magic packet to a pre-configured (or an arbitrary) machine in the local network\n";
    usage += &ssh_unlock_disk_usage(config);
    usage += " - log into a pre-configured machine via SSH (as root) and send the password to unlock a LUKS-encrypted root partition\n";
    usage += &ssh_exec_usage(config);
    usage += " - log into a machine via SSH and execute one command. Waits until the process finishes or forcefully exits after 20 seconds\n";

    usage += "\nCommands from plugins:\n";
    for cmd in ctx_lock.plugin_commands.values() {
        usage += &format!(
            "{} (from {}) - {}\n",
            &cmd.usage, &cmd.plugin_id, &cmd.description
        );
    }

    usage += "\nDeltaChat Remote Control Bot by slavasil, 2026\n";

    usage
}

pub async fn auth_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let mut ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !AUTH_REQUIRED {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Authentication is disabled".to_owned(),
        )
        .await?;
        return Ok(());
    }

    if ctx_lock.authed_contacts.contains(&contact_id) {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Already authenticated".to_owned()).await?;
        return Ok(());
    }

    let Some(password) = args.first() else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Usage: /auth <password>".to_owned(),
        )
        .await?;
        return Ok(());
    };

    if *password == ctx_lock.config.auth.password {
        ctx_lock.authed_contacts.insert(contact_id);
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            "Authentication successful!".to_owned(),
        )
        .await?;
    } else {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, "Incorrect password".to_owned()).await?;
    }

    Ok(())
}

pub async fn wol_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !ensure_auth(&dchat_ctx_lock, &ctx_lock, chat_id, contact_id).await? {
        return Ok(());
    }

    let mac = match args.first() {
        None => match ctx_lock.config.machines.values().find(|m| m.is_default) {
            Some(machine) => machine.mac,
            None => {
                let usage = wol_usage(&ctx_lock.config);
                chat::send_text_msg(&dchat_ctx_lock, chat_id, usage).await?;
                return Ok(());
            }
        },
        Some(arg) => {
            if let Some(machine) = ctx_lock.config.machines.get(*arg) {
                machine.mac
            } else if let Ok(mac) = MacAddress::parse_str(arg) {
                mac
            } else {
                let usage = wol_usage(&ctx_lock.config);
                chat::send_text_msg(&dchat_ctx_lock, chat_id, usage).await?;
                return Ok(());
            }
        }
    };

    let reply = match send_magic_packet(mac).await {
        Ok(()) => format!("Magic packet sent to {}.", mac_to_string(mac)),
        Err(e) => format!("Failed to send magic packet: {e:#}"),
    };
    chat::send_text_msg(&dchat_ctx_lock, chat_id, reply).await?;

    Ok(())
}

fn wol_usage(config: &BotConfig) -> String {
    let mut forms: Vec<String> = Vec::new();

    if config.machines.values().any(|m| m.is_default) {
        forms.push("/wol".to_owned());
    }

    let mut names: Vec<&str> = config.machines.keys().map(String::as_str).collect();
    names.sort();
    if !names.is_empty() {
        forms.push(format!("/wol <{}>", names.join("|")));
    }

    forms.push("/wol <XX:XX:XX:XX:XX:XX>".to_owned());

    format!("Usage: {}", forms.join(" OR "))
}

fn mac_to_string(mac: MacAddress) -> String {
    let b = mac.as_bytes();
    format!(
        "{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}",
        b[0], b[1], b[2], b[3], b[4], b[5]
    )
}

async fn send_magic_packet(mac: MacAddress) -> AnyhowResult<()> {
    let mut packet = [0u8; 102];
    packet[..6].fill(0xFF);
    let mac_bytes = mac.as_bytes();
    for i in 0..16 {
        packet[6 + i * 6..6 + (i + 1) * 6].copy_from_slice(mac_bytes);
    }

    let socket = tokio::net::UdpSocket::bind("0.0.0.0:0")
        .await
        .context("Failed to bind UDP socket for WoL")?;
    socket
        .set_broadcast(true)
        .context("Failed to enable broadcast on UDP socket")?;
    socket
        .send_to(&packet, "255.255.255.255:9")
        .await
        .context("Failed to send magic packet")?;

    Ok(())
}

fn ssh_unlock_disk_usage(config: &BotConfig) -> String {
    if !config.machines.is_empty() {
        let list = config
            .machines
            .keys()
            .cloned()
            .collect::<Vec<_>>()
            .join("|");
        if config.default_machine().is_some() {
            format!("Usage: /ssh-unlock-disk [{list}] <password>")
        } else {
            format!("Usage: /ssh-unlock-disk {list} <password>")
        }
    } else {
        "Command /ssh-unlock-disk unavailable: no machines configured".to_string()
    }
}

pub async fn ssh_unlock_disk_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !ensure_auth(&dchat_ctx_lock, &ctx_lock, chat_id, contact_id).await? {
        return Ok(());
    }

    if args.is_empty() || args.len() > 2 {
        log::warn!("wrong number of args");
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            ssh_unlock_disk_usage(&ctx_lock.config),
        )
        .await?;
        return Ok(());
    }

    let Some((machine_name, machine)) = ({
        if args.len() > 1 {
            ctx_lock.config.machines.get(args[0]).map(|m| (args[0], m))
        } else if let Some((name, machine)) = ctx_lock.config.default_machine() {
            log::info!("default machine: {:?} {:?}", name, machine);
            Some((name, machine))
        } else {
            log::warn!("no machine found :/");
            None
        }
    }) else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            ssh_unlock_disk_usage(&ctx_lock.config),
        )
        .await?;
        return Ok(());
    };

    let Some(remote_unlock_cfg) = machine.remote_unlock.as_ref() else {
        chat::send_text_msg(
            &dchat_ctx_lock,
            chat_id,
            format!("Remote unlock is not configured for machine '{machine_name}"),
        )
        .await?;
        return Ok(());
    };

    let ssh_key_relative_path: PathBuf = remote_unlock_cfg.ssh_key_file.clone().into();
    let ssh_key_path = data_path().join(ssh_key_relative_path);
    log::debug!("Reading SSH key from {:?}", &ssh_key_path);
    let mut ssh_key_file = File::open(ssh_key_path)?;
    let mut ssh_key = String::new();
    ssh_key_file.read_to_string(&mut ssh_key)?;
    drop(ssh_key_file);

    let password = args[args.len() - 1].to_owned();
    let host = machine.static_ip.to_string();
    let ssh_key =
        russh::keys::PrivateKey::from_openssh(ssh_key).context("Invalid remote unlock SSH key")?;

    drop(ctx_lock);
    drop(dchat_ctx_lock);

    tokio::spawn(async move {
        let ssh_config = Arc::new(russh::client::Config {
            inactivity_timeout: None,
            ..Default::default()
        });
        // TODO implement real host key verification (needs config file changes)
        const SSH_TIMEOUT: Duration = Duration::from_secs(15);
        log::info!("Connecting to SSH {}:22", &host);
        let mut connection = timeout(
            SSH_TIMEOUT,
            russh::client::connect(ssh_config, (host, 22), crate::ssh::ClientHandler {}),
        )
        .await
        .context("Connection to host timed out")?
        .context("Cannot connect to host")?;

        log::info!("Trying to authenticate as root");
        let auth_result = timeout(
            SSH_TIMEOUT,
            connection.authenticate_publickey(
                "root",
                PrivateKeyWithHashAlg::new(
                    Arc::new(ssh_key),
                    connection.best_supported_rsa_hash().await?.flatten(),
                ),
            ),
        )
        .await
        .context("SSH authentication timed out")?
        .context("SSH auth against remote unlock target failed")?;

        // ignoring both fields because we only have the key
        if let AuthResult::Failure {
            remaining_methods: _,
            partial_success: _,
        } = auth_result
        {
            anyhow::bail!("SSH authentication failed");
        }

        log::info!("Opening session");
        let mut channel = timeout(SSH_TIMEOUT, connection.channel_open_session())
            .await
            .context("Timed out")?
            .context("Cannot open SSH session")?;

        log::info!("Requesting shell");
        timeout(SSH_TIMEOUT, channel.request_shell(true))
            .await
            .context("Timed out")?
            .context("Cannot request disk password prompt")?;

        log::info!("Sending password");
        let input = password + "\n";
        timeout(SSH_TIMEOUT, channel.data(input.as_bytes()))
            .await
            .context("Timed out")?
            .context("Cannot send password via SSH")?;

        loop {
            let Some(msg) = timeout(SSH_TIMEOUT, channel.wait())
                .await
                .context("Timed out")?
            else {
                break;
            };
            log::info!("Received message from SSH server: {:?}", msg);
        }

        Ok(())
    })
    .await?
    .context("SSH unlock error")?;

    let dchat_ctx_lock = dchat_ctx.lock().await;
    chat::send_text_msg(&dchat_ctx_lock, chat_id, "Unlock password sent".to_owned()).await?;

    Ok(())
}

fn ssh_exec_usage(config: &BotConfig) -> String {
    if !config.machines.is_empty() {
        let list = config
            .machines
            .keys()
            .cloned()
            .collect::<Vec<_>>()
            .join("|");
        format!("Usage: /ssh-exec <username>@({list}|<hostname>|<IP address>) <command...>")
    } else {
        "Command /ssh-exec unavailable: no machines configured".to_string()
    }
}

pub async fn ssh_exec_command(
    dchat_ctx: Arc<Mutex<Context>>,
    ctx: Arc<Mutex<BotContext>>,
    msg: Message,
    args: &[&str],
) -> AnyhowResult<()> {
    let chat_id = msg.get_chat_id();
    let contact_id = msg.get_from_id();

    let ctx_lock = ctx.lock().await;
    let dchat_ctx_lock = dchat_ctx.lock().await;

    if !ensure_auth(&dchat_ctx_lock, &ctx_lock, chat_id, contact_id).await? {
        return Ok(());
    }

    let cmd = Command::new("ssh-exec").args([
        Arg::new("password").short('P').long("password"),
        Arg::new("target"),
        Arg::new("command")
            .action(ArgAction::Append)
            .num_args(1..)
            .trailing_var_arg(true),
    ]);
    let mut args_with_dummy_program_name = vec!["ssh-exec"];
    args_with_dummy_program_name.extend_from_slice(args);
    let Ok(cmdline_opts) = cmd.try_get_matches_from(args_with_dummy_program_name) else {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, ssh_exec_usage(&ctx_lock.config)).await?;
        return Ok(());
    };

    let (Some(target), Some(command)) = (
        cmdline_opts.get_one::<String>("target"),
        cmdline_opts.get_many::<String>("command"),
    ) else {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, ssh_exec_usage(&ctx_lock.config)).await?;
        return Ok(());
    };

    if !target.contains('@') {
        chat::send_text_msg(&dchat_ctx_lock, chat_id, ssh_exec_usage(&ctx_lock.config)).await?;
        return Ok(());
    }

    let command = command.fold(String::new(), |mut out, part| {
        if !out.is_empty() {
            out.push(' ');
        }
        out.push_str(part);
        out
    });

    let (username, host) = target.split_once('@').unwrap();
    let username = username.to_owned();
    let (host_addr, key_filename) = match ctx_lock.config.machines.get(host) {
        Some(machine) => (
            append_port_if_needed(&machine.static_ip.to_string(), 22),
            Some(machine.remote_access.ssh_key_file.clone()),
        ),
        None => (append_port_if_needed(host, 22), None),
    };

    let key = key_filename
        .map(|filename| -> AnyhowResult<String> {
            let filename = data_path().join(filename);
            log::debug!("Reading SSH key from {:?}", &filename);
            let mut ssh_key_file = File::open(filename)?;
            let mut ssh_key = String::new();
            ssh_key_file.read_to_string(&mut ssh_key)?;
            Ok(ssh_key)
        })
        .transpose()
        .context("Cannot read the SSH key file")?;

    let password = cmdline_opts.get_one::<String>("password").cloned();

    let report_message_id = chat::send_text_msg(
        &dchat_ctx_lock,
        chat_id,
        format!("Executing {} on {}@{}...", &command, &username, &host_addr),
    )
    .await
    .context("Cannot send reply")?;

    drop(ctx_lock);
    drop(dchat_ctx_lock);

    tokio::spawn(async move {
        const TIMEOUT: Duration = Duration::from_secs(15);
        const EXIT_TIMEOUT: Duration = Duration::from_secs(20);

        let connection = ssh::connect_ssh(TIMEOUT, TIMEOUT, &host_addr, &username, key, password)
            .await
            .context(format!("Cannot connect to {}", &host_addr))?;

        log::info!("Opening SSH channel");
        let mut session = timeout(TIMEOUT, connection.channel_open_session())
            .await
            .context("Open SSH channel timed out")?
            .context("Cannot open SSH channel")?;

        log::info!("Executing SSH command \"{}\"", &command);
        timeout(TIMEOUT, session.exec(false, command.as_bytes()))
            .await
            .context("SSH exec timed out")?
            .context("Cannot execute command via SSH")?;

        let exit_deadline = Instant::now().add(EXIT_TIMEOUT);
        let mut report_text = String::new();
        loop {
            match timeout_at(exit_deadline, session.wait()).await {
                Ok(Some(msg)) => match msg {
                    russh::ChannelMsg::Data { data } => {
                        let dchat_ctx_lock = dchat_ctx.lock().await;
                        report_text += &String::from_utf8_lossy(&data);
                        chat::send_edit_request(
                            &dchat_ctx_lock,
                            report_message_id,
                            report_text.clone(),
                        )
                        .await
                        .context("Cannot edit message")?;
                    }
                    russh::ChannelMsg::ExtendedData { data, ext: 1 } => {
                        let dchat_ctx_lock = dchat_ctx.lock().await;
                        report_text += &String::from_utf8_lossy(&data);
                        chat::send_edit_request(
                            &dchat_ctx_lock,
                            report_message_id,
                            report_text.clone(),
                        )
                        .await
                        .context("Cannot edit message")?;
                    }
                    _ => {
                        log::info!("Message from SSH channel: {msg:?}");
                    }
                },
                Ok(None) => {
                    log::info!("No message from SSH channel. Is this EOF?");
                    break;
                }
                Err(_) => {
                    log::error!("SSH exec command timed out");

                    let dchat_ctx_lock = dchat_ctx.lock().await;
                    report_text += "\n\nRead timed out";
                    chat::send_edit_request(
                        &dchat_ctx_lock,
                        report_message_id,
                        report_text.clone(),
                    )
                    .await
                    .context("Cannot edit message")?;

                    break;
                }
            }
        }
        AnyhowResult::<()>::Ok(())
    })
    .await?
    .context("SSH exec error")?;

    Ok(())
}

async fn ensure_auth(
    dchat_ctx: &Context,
    ctx: &BotContext,
    chat_id: ChatId,
    contact_id: ContactId,
) -> AnyhowResult<bool> {
    if !AUTH_REQUIRED {
        return Ok(true);
    }
    if !ctx.authed_contacts.contains(&contact_id) {
        chat::send_text_msg(
            dchat_ctx,
            chat_id,
            "Authenticate yourself first with '/auth <password>'.".to_owned(),
        )
        .await?;
        Ok(false)
    } else {
        Ok(true)
    }
}