//! `dragoon-ctl` — controller CLI for remote-executor.

use std::io::{BufRead, IsTerminal};
use std::path::PathBuf;
use std::time::Duration;

use anyhow::{anyhow, Context, Result};
use base64::{engine::general_purpose::STANDARD as B64, Engine};
use clap::{Parser, Subcommand};
use dragoon_proto::{
    models::{TaskKind, TaskLimits},
    signer::{agent::AgentSigner, keyfile::KeyFileSigner, Signer},
};
use dragoon_ctl::client::CtlClient;
use serde_json::Value;

#[derive(Parser, Debug)]
#[command(name = "dragoon-ctl", version, about = "remote-executor controller")]
struct Cli {
    #[command(subcommand)]
    command: Cmd,
}

#[derive(Subcommand, Debug)]
enum Cmd {
    Login {
        #[arg(long)]
        server: String,
        #[arg(long, short = 'u')]
        username: String,
        #[arg(long)]
        ssh_key: Option<PathBuf>,
        #[arg(long)]
        ssh_fingerprint: Option<String>,
        #[arg(long)]
        no_verify_tls: bool,
    },
    Logout,
    Worker {
        #[command(subcommand)]
        cmd: WorkerCmd,
    },
    Submit {
        worker: String,
        #[arg(short = 'c')]
        cmd: Option<String>,
        #[arg(short = 'f')]
        file: Option<PathBuf>,
        #[arg(long)]
        collect: Option<String>,
        #[arg(long)]
        timeout: Option<u64>,
        #[arg(long)]
        max_size: Option<u64>,
        #[arg(long)]
        max_files: Option<u64>,
    },
    Status {
        task_id: String,
    },
    Cancel {
        task_id: String,
    },
    Tail {
        task_id: String,
        #[arg(long, default_value = "stdout")]
        stream: String,
    },
    Pull {
        task_id: String,
        #[arg(long, default_value = ".")]
        out: PathBuf,
    },
    Fetch {
        worker: String,
        #[arg(long)]
        path: String,
        #[arg(long)]
        glob: String,
    },
    Events {
        #[arg(long, default_value_t = 0)]
        since: i64,
        #[arg(long)]
        watch: bool,
        #[arg(long)]
        all: bool,
        #[arg(long)]
        no_ack: bool,
    },
}

#[derive(Subcommand, Debug)]
enum WorkerCmd {
    List,
}

fn config_dir() -> PathBuf {
    if let Ok(d) = std::env::var("REMOTE_EXECUTOR_CONFIG_DIR") {
        return PathBuf::from(d);
    }
    let home = std::env::var_os("HOME").unwrap_or_else(|| "/tmp".into());
    PathBuf::from(home).join(".config").join("remote-executor")
}

fn session_path() -> PathBuf {
    config_dir().join("session.json")
}

fn save_session(server_url: &str, token: &str, fingerprint: &str, expires_at: &str, ssh_key: Option<&PathBuf>) -> Result<()> {
    let p = session_path();
    if let Some(parent) = p.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let mut payload = serde_json::json!({
        "server_url": server_url,
        "session_token": token,
        "fingerprint": fingerprint,
        "expires_at": expires_at,
    });
    if let Some(k) = ssh_key {
        payload["ssh_key"] = Value::String(k.to_string_lossy().to_string());
    }
    std::fs::write(&p, serde_json::to_string_pretty(&payload)?)?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let _ = std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600));
    }
    Ok(())
}

fn load_session() -> Result<Option<Value>> {
    let p = session_path();
    if !p.exists() {
        return Ok(None);
    }
    Ok(Some(serde_json::from_str(&std::fs::read_to_string(&p)?)?))
}

fn clear_session() {
    let _ = std::fs::remove_file(session_path());
}

fn build_signer(
    ssh_key: Option<&PathBuf>,
    ssh_fingerprint: Option<&str>,
) -> Result<Box<dyn Signer + Send + Sync>> {
    if let Some(p) = ssh_key {
        Ok(Box::new(KeyFileSigner::from_openssh_file(p, None)?))
    } else {
        Ok(Box::new(AgentSigner::from_env(ssh_fingerprint, None)?))
    }
}

fn client_from_session() -> Result<CtlClient> {
    let sess = load_session()?
        .ok_or_else(|| anyhow!("not logged in; run 'ctl login' first"))?;
    let server_url = sess["server_url"].as_str().unwrap_or("");
    let token = sess["session_token"].as_str().map(String::from);
    let signer: Box<dyn Signer + Send + Sync> = if let Some(k) = sess.get("ssh_key").and_then(Value::as_str) {
        Box::new(KeyFileSigner::from_openssh_file(PathBuf::from(k), None)?)
    } else {
        let fp = sess.get("fingerprint").and_then(Value::as_str);
        Box::new(AgentSigner::from_env(fp, None)?)
    };
    CtlClient::new(server_url, signer, true, token)
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
        )
        .init();

    match cli.command {
        Cmd::Login { server, username, ssh_key, ssh_fingerprint, no_verify_tls } => {
            let signer = build_signer(ssh_key.as_ref(), ssh_fingerprint.as_deref())?;
            let fp = signer.fingerprint().to_string();
            let mut client = CtlClient::new(&server, signer, !no_verify_tls, None)?;
            let interactive = std::io::stdin().is_terminal();
            let read_line = |prompt: &str| -> Result<String> {
                if interactive {
                    Ok(rpassword::prompt_password(prompt)?)
                } else {
                    let stdin = std::io::stdin();
                    let mut line = String::new();
                    stdin.lock().read_line(&mut line)?;
                    Ok(line.trim_end_matches(['\n', '\r']).to_string())
                }
            };
            let pw = read_line("Password: ")?;
            let totp = read_line("TOTP code: ")?;
            let resp = client.login(&username, &pw, totp.trim()).await?;
            save_session(
                &server,
                &resp.session_token,
                &fp,
                &resp.expires_at.to_rfc3339(),
                ssh_key.as_ref(),
            )?;
            println!("logged in as {username} until {}", resp.expires_at);
            Ok(())
        }
        Cmd::Logout => {
            if let Ok(mut c) = client_from_session() {
                let _ = c.logout().await;
            }
            clear_session();
            println!("logged out");
            Ok(())
        }
        Cmd::Worker { cmd } => match cmd {
            WorkerCmd::List => {
                let c = client_from_session()?;
                for w in c.list_workers().await? {
                    println!(
                        "{:<20} {:<8} pwd={} task={}",
                        w["name"].as_str().unwrap_or("-"),
                        w["status"].as_str().unwrap_or("-"),
                        w["current_pwd"].as_str().unwrap_or("-"),
                        w["current_task_id"].as_str().unwrap_or("-"),
                    );
                }
                Ok(())
            }
        },
        Cmd::Submit { worker, cmd, file, collect, timeout, max_size, max_files } => {
            if cmd.is_some() == file.is_some() {
                anyhow::bail!("exactly one of -c or -f is required");
            }
            let (kind, payload) = if let Some(c) = cmd {
                (TaskKind::Command, c)
            } else {
                let f = file.unwrap();
                let s = std::fs::read_to_string(&f)
                    .with_context(|| format!("read {}", f.display()))?;
                (TaskKind::Script, s)
            };
            let limits = TaskLimits {
                max_total_size_bytes: max_size.unwrap_or(100 * 1024 * 1024),
                max_files: max_files.unwrap_or(1000),
                timeout_sec: timeout,
            };
            let collect_globs = collect
                .map(|s| s.split(',').map(|x| x.trim().to_string()).collect::<Vec<_>>())
                .unwrap_or_default();
            let c = client_from_session()?;
            let t = c.submit(&worker, kind, &payload, collect_globs, limits, None).await?;
            println!("{}", t.task_id);
            Ok(())
        }
        Cmd::Status { task_id } => {
            let c = client_from_session()?;
            let t = c.get_task(&task_id).await?;
            println!("{}", serde_json::to_string_pretty(&t)?);
            Ok(())
        }
        Cmd::Cancel { task_id } => {
            let c = client_from_session()?;
            let t = c.cancel_task(&task_id).await?;
            println!("{:?}", t.state);
            Ok(())
        }
        Cmd::Tail { task_id, stream } => {
            let c = client_from_session()?;
            let mut seq: u64 = 0;
            loop {
                let v = c.tail_log(&task_id, seq, &stream).await?;
                let data_b64 = v["data_b64"].as_str().unwrap_or("");
                if !data_b64.is_empty() {
                    let bytes = B64.decode(data_b64.as_bytes()).unwrap_or_default();
                    print!("{}", String::from_utf8_lossy(&bytes));
                    seq = v["seq"].as_u64().unwrap_or(seq);
                }
                if v["eof"].as_bool().unwrap_or(false) {
                    return Ok(());
                }
                if data_b64.is_empty() {
                    tokio::time::sleep(Duration::from_millis(500)).await;
                }
            }
        }
        Cmd::Pull { task_id, out } => {
            let c = client_from_session()?;
            let t = c.get_task(&task_id).await?;
            std::fs::create_dir_all(&out)?;
            for art in &t.artifacts {
                let bytes = c.download_artifact(&task_id, &art.path).await?;
                let dest = out.join(&art.path);
                if let Some(p) = dest.parent() {
                    std::fs::create_dir_all(p)?;
                }
                std::fs::write(&dest, &bytes)?;
                println!("{} ({} bytes)", dest.display(), bytes.len());
            }
            Ok(())
        }
        Cmd::Fetch { worker, path, glob } => {
            let c = client_from_session()?;
            let t = c.fetch(&worker, &path, &glob).await?;
            println!("{}", t.task_id);
            Ok(())
        }
        Cmd::Events { since, watch, all, no_ack } => {
            let c = client_from_session()?;
            let mut cursor = since;
            loop {
                let v = c.fetch_messages(cursor, all, watch).await?;
                let msgs = v["messages"].as_array().cloned().unwrap_or_default();
                let mut new_ids: Vec<i64> = Vec::new();
                for m in &msgs {
                    let id = m["id"].as_i64().unwrap_or(0);
                    let kind = m["kind"].as_str().unwrap_or("?");
                    let task = m["task_id"].as_str().unwrap_or("-");
                    let worker = m["worker_name"].as_str().unwrap_or("-");
                    let state = m["payload"]["state"].as_str().unwrap_or("");
                    let err = m["error"].as_str().unwrap_or("");
                    let cmd = m["payload"]["task_payload"].as_str().unwrap_or("");
                    let exit = m["payload"]["exit_code"].as_i64();
                    let ts = m["created_at"].as_str().unwrap_or("");
                    println!(
                        "#{:>5}  {}  {:<14}  worker={}  task={}  state={}",
                        id, ts, kind, worker, task, state
                    );
                    if !err.is_empty() {
                        println!("        error: {err}");
                    }
                    if !cmd.is_empty() {
                        let preview = if cmd.len() <= 200 { cmd.to_string() } else { format!("{}...", &cmd[..200]) };
                        println!("        cmd:   {preview}");
                    }
                    if let Some(e) = exit {
                        println!("        exit:  {e}");
                    }
                    new_ids.push(id);
                }
                if let Some(next) = v.get("next_since_id").and_then(Value::as_i64) {
                    if next > cursor {
                        cursor = next;
                    }
                }
                if !new_ids.is_empty() && !no_ack {
                    let _ = c.ack_messages(new_ids).await;
                }
                if !watch {
                    return Ok(());
                }
                if msgs.is_empty() {
                    tokio::time::sleep(Duration::from_millis(500)).await;
                }
            }
        }
    }
}
