use anyhow::{Context, Result};
use russh::{ChannelMsg, Disconnect};
use tokio::io::{AsyncWriteExt, Stderr, Stdout};
use tracing;

use super::connect::ClientHandler;

/// Execute a command and write output to stdout/stderr.
pub(crate) async fn exec_on_session(
    session: &russh::client::Handle<ClientHandler>,
    cmd: &str,
    stdout: &mut Stdout,
    stderr: &mut Stderr,
) -> Result<u32> {
    let (exit_code, stdout_buf, stderr_buf) = exec_collect_inner(session, cmd).await?;
    stdout.write_all(&stdout_buf).await?;
    stdout.flush().await?;
    stderr.write_all(&stderr_buf).await?;
    stderr.flush().await?;
    Ok(exit_code)
}

/// Execute a command and collect stdout/stderr into buffers.
pub(crate) async fn exec_collect(
    session: &russh::client::Handle<ClientHandler>,
    cmd: &str,
) -> Result<(u32, Vec<u8>, Vec<u8>)> {
    exec_collect_inner(session, cmd).await
}

/// Core execution logic: open channel, execute command, collect output.
async fn exec_collect_inner(
    session: &russh::client::Handle<ClientHandler>,
    cmd: &str,
) -> Result<(u32, Vec<u8>, Vec<u8>)> {
    let start = std::time::Instant::now();
    tracing::debug!(command = cmd, "exec command start");
    let mut channel = session
        .channel_open_session()
        .await
        .with_context(|| "open session channel failed")
        .map_err(|e| {
            tracing::error!(command = cmd, error = %e, "exec command failed");
            e
        })?;
    channel
        .exec(true, cmd)
        .await
        .with_context(|| "exec failed")
        .map_err(|e| {
            tracing::error!(command = cmd, error = %e, "exec command failed");
            e
        })?;

    let mut exit_code: Option<u32> = None;
    let mut stdout_buf: Vec<u8> = Vec::new();
    let mut stderr_buf: Vec<u8> = Vec::new();

    loop {
        let Some(msg) = channel.wait().await else {
            break;
        };
        match msg {
            ChannelMsg::Data { ref data } => stdout_buf.extend_from_slice(data),
            ChannelMsg::ExtendedData { ref data, .. } => stderr_buf.extend_from_slice(data),
            ChannelMsg::ExitStatus { exit_status } => exit_code = Some(exit_status),
            _ => {}
        }
    }

    let exit_code = exit_code.unwrap_or(255);
    tracing::info!(exit_code = exit_code, elapsed = ?start.elapsed(), "exec command done");
    Ok((exit_code, stdout_buf, stderr_buf))
}

pub(crate) async fn disconnect_session(
    session: &russh::client::Handle<ClientHandler>,
) -> Result<()> {
    session
        .disconnect(Disconnect::ByApplication, "", "en")
        .await
        .ok();
    Ok(())
}
