use anyhow::Result;
use clap::Parser;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::broadcast;

mod cli;

use cli::{Cli, Commands};
use sshr::ssh::{
    AuthMethod, ExecService, PtyOpenRequest, PtyService, RusshTerminalService, RusshTool,
    SessionEvent, SessionId, SshConfig,
};

#[derive(Clone)]
struct TerminalApp {
    service: RusshTerminalService,
}

#[tokio::main]
async fn main() -> Result<()> {
    dotenvy::dotenv().ok();
    let cli = Cli::parse();

    let tool = RusshTool::new();
    match cli.command {
        Commands::Validate(args) => {
            ExecService::validate(
                &tool,
                config_from_args(
                    args.host,
                    args.port,
                    args.user,
                    args.password,
                    args.key_path,
                    args.key_passphrase,
                    args.timeout_ms,
                ),
            )
            .await?;
            println!("ok");
        }
        Commands::Exec(args) => {
            let result = ExecService::exec(
                &tool,
                config_from_args(
                    args.host,
                    args.port,
                    args.user,
                    args.password,
                    args.key_path,
                    args.key_passphrase,
                    args.timeout_ms,
                ),
                sshr::ssh::ExecRequest::new(args.command.join(" ")),
            )
            .await?;
            tokio::io::stdout().write_all(&result.stdout).await?;
            tokio::io::stdout().flush().await?;
            tokio::io::stderr().write_all(&result.stderr).await?;
            tokio::io::stderr().flush().await?;
            std::process::exit(result.exit_code as i32);
        }
        Commands::Shell(args) => {
            ExecService::open_shell(
                &tool,
                config_from_args(
                    args.host,
                    args.port,
                    args.user,
                    args.password,
                    args.key_path,
                    args.key_passphrase,
                    args.timeout_ms,
                ),
            )
            .await?;
        }
        Commands::Pty(args) => {
            let app = TerminalApp {
                service: RusshTerminalService::new(),
            };
            run_terminal_demo(
                &app,
                config_from_args(
                    args.host,
                    args.port,
                    args.user,
                    args.password,
                    args.key_path,
                    args.key_passphrase,
                    args.timeout_ms,
                ),
            )
            .await?;
        }
    }
    Ok(())
}

fn config_from_args(
    host: String,
    port: u16,
    user: String,
    password: Option<String>,
    key_path: Option<String>,
    key_passphrase: Option<String>,
    timeout_ms: u64,
) -> SshConfig {
    let auth = if let Some(path) = key_path {
        AuthMethod::PrivateKey {
            private_key_path: path,
            passphrase: key_passphrase,
        }
    } else {
        AuthMethod::Password {
            password: password.unwrap_or_default(),
        }
    };

    SshConfig {
        host,
        port,
        user,
        auth,
        timeout_ms,
        host_key_policy: Default::default(),
        known_hosts_path: None,
    }
}

async fn run_terminal_demo(app: &TerminalApp, cfg: SshConfig) -> Result<()> {
    let (cols, rows) = crossterm::terminal::size().unwrap_or((80, 24));
    let session_id = app
        .service
        .open(
            cfg,
            PtyOpenRequest {
                term: std::env::var("TERM").unwrap_or_else(|_| "xterm-256color".to_string()),
                cols,
                rows,
                ..PtyOpenRequest::default()
            },
        )
        .await?;
    let mut events = app.service.subscribe(session_id).await?;

    crossterm::terminal::enable_raw_mode()?;
    let raw_mode_guard = RawModeGuard;

    let app_clone = app.clone();
    let stdin_task = tokio::spawn(async move { forward_stdin(app_clone, session_id).await });

    let mut stdout = tokio::io::stdout();
    let mut stderr = tokio::io::stderr();

    loop {
        match events.recv().await {
            Ok(SessionEvent::Stdout(data)) => {
                stdout.write_all(&data).await?;
                stdout.flush().await?;
            }
            Ok(SessionEvent::Stderr(data)) => {
                stderr.write_all(&data).await?;
                stderr.flush().await?;
            }
            Ok(SessionEvent::Error(message)) => {
                stderr.write_all(message.as_bytes()).await?;
                stderr.write_all(b"\n").await?;
                stderr.flush().await?;
            }
            Ok(SessionEvent::Closed { exit_status }) => {
                if let Some(code) = exit_status {
                    stderr
                        .write_all(format!("\n[session closed: {code}]\n").as_bytes())
                        .await?;
                    stderr.flush().await?;
                }
                break;
            }
            Err(broadcast::error::RecvError::Closed) => break,
            Err(broadcast::error::RecvError::Lagged(_)) => continue,
        }
    }

    stdin_task.abort();
    let _ = stdin_task.await;
    app.service.close(session_id).await.ok();
    drop(raw_mode_guard);
    Ok(())
}

async fn forward_stdin(app: TerminalApp, session_id: SessionId) {
    let mut stdin = tokio::io::stdin();
    let mut buf = [0u8; 4096];

    loop {
        let n = match stdin.read(&mut buf).await {
            Ok(0) => break,
            Ok(n) => n,
            Err(_) => break,
        };

        if app
            .service
            .write_input(session_id, buf[..n].to_vec())
            .await
            .is_err()
        {
            break;
        }
    }
}

struct RawModeGuard;

impl Drop for RawModeGuard {
    fn drop(&mut self) {
        crossterm::terminal::disable_raw_mode().ok();
    }
}
