forklaunch 1.20.1

#[cfg(not(unix))]
use std::fs::create_dir_all;
#[cfg(unix)]
use std::os::unix::fs::{DirBuilderExt, OpenOptionsExt};
use std::{
    fs::OpenOptions,
    io::{IsTerminal, Write},
    thread::sleep,
    time::Duration,
};

use anyhow::{Result, bail};
use clap::{Arg, ArgMatches, Command};
use serde::{Deserialize, Serialize};
use termcolor::{Color, ColorChoice, StandardStream, WriteColor};

use crate::{
    CliCommand,
    constants::get_iam_api_url,
    core::{
        command::command,
        token::{API_KEY_PREFIX, exchange_api_key, get_token_path},
    },
};

pub(super) struct LoginCommand;

impl LoginCommand {
    pub(super) fn new() -> Self {
        Self {}
    }
}

#[derive(Debug, Deserialize)]
struct DeviceCodeResponse {
    device_code: String,
    user_code: String,
    verification_uri: String,
    verification_uri_complete: Option<String>,
    interval: Option<i64>,
}

#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
}

#[derive(Debug, Deserialize)]
struct TokenErrorResponse {
    error: String,
    error_description: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
struct TokenData {
    access_token: String,
    refresh_token: String,
    expires_at: i64,
}

/// Login with a credential rather than a browser — the path for CI, a
/// container, and an agent operating the platform unattended.
///
/// Two kinds of value arrive here. An **API key** (`flk_…`, issued by
/// `POST /iam/service-accounts`) is a long-lived machine credential: it is
/// exchanged for a JWT now and kept, so the session renews itself when that
/// JWT runs out. That renewal is what makes unattended operation actually
/// unattended — a device-flow session eventually needs a person at a
/// browser, and a bare JWT simply stops working. A **raw JWT** is still
/// accepted, for the case where something upstream already minted one.
pub fn login_with_token(api_token: &str) -> Result<()> {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);
    let is_api_key = api_token.starts_with(API_KEY_PREFIX);

    log_info!(stdout, "Forklaunch CLI Login (API Token)");

    let token_storage = if is_api_key {
        log_info!(stdout, "Exchanging API key for a session...");
        let (access_token, expires_at) = exchange_api_key(api_token)?;
        TokenData {
            access_token,
            // The key is what renews the session, so it is stored where the
            // refresh path already looks.
            refresh_token: api_token.to_string(),
            expires_at,
        }
    } else {
        log_info!(stdout, "Validating API token...");
        TokenData {
            access_token: api_token.to_string(),
            refresh_token: String::new(),
            // Read the real expiry from the token when it has one. Recording
            // "never" meant the CLI kept presenting an expired token until the
            // server's 401 wiped the login file.
            expires_at: crate::core::token::jwt_expiry(api_token).unwrap_or(i64::MAX),
        }
    };

    let token_path = get_token_path()?;

    // Ensure parent directory exists with owner-only permissions (0o700)
    if let Some(parent) = token_path.parent() {
        #[cfg(unix)]
        {
            use std::fs::DirBuilder;
            let mut builder = DirBuilder::new();
            builder.recursive(true);
            builder.mode(0o700);
            builder.create(parent)?;
        }

        #[cfg(not(unix))]
        {
            create_dir_all(parent)?;
        }
    }

    let toml_content = toml::to_string(&token_storage)?;

    // Write token file with owner-only permissions (0o600)
    #[cfg(unix)]
    {
        use std::io::Write as IoWrite;

        let mut file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o600)
            .open(&token_path)?;

        file.write_all(toml_content.as_bytes())?;
    }

    #[cfg(not(unix))]
    {
        use std::fs::write;
        write(&token_path, toml_content)?;
    }

    writeln!(stdout)?;
    log_header!(
        stdout,
        Color::Green,
        "Successfully logged in with API token!"
    );
    if is_api_key {
        writeln!(
            stdout,
            "This session renews itself from the key, so unattended runs keep working. \
             Revoke the key from the platform if it leaks."
        )?;
    } else {
        writeln!(
            stdout,
            "This is a raw token and cannot be renewed; when it expires, log in again. \
             An API key (flk_...) from a service account renews itself instead."
        )?;
    }

    Ok(())
}

/// Is there a human at this terminal to read a code and open a browser?
///
/// The device flow asks the user to visit a URL and type a code, then polls
/// until they do or the code expires ten minutes later. With nobody there —
/// CI, an agent, a container — that ten minutes is pure waiting, and the run
/// ends with "The device code has expired" instead of saying what was
/// actually wrong. `FORKLAUNCH_FORCE_DEVICE_LOGIN=1` overrides, for the rare
/// case of a browser but no tty.
fn device_login_decision(forced: bool, stdin_tty: bool, stdout_tty: bool) -> bool {
    forced || (stdin_tty && stdout_tty)
}

fn device_login_is_usable() -> bool {
    device_login_decision(
        std::env::var("FORKLAUNCH_FORCE_DEVICE_LOGIN").is_ok_and(|v| !v.is_empty()),
        std::io::stdin().is_terminal(),
        std::io::stdout().is_terminal(),
    )
}

/// Interactive device flow login (default)
pub fn login() -> Result<()> {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);
    let api_url = get_iam_api_url();

    if !device_login_is_usable() {
        bail!(
            "`forklaunch login` needs a terminal: it prints a code for you to enter in a browser, \
             and with nobody there it polls for ten minutes and then fails. For CI, an agent or a \
             container, authenticate headlessly instead:\n  \
             forklaunch login --token <api-token>   (or set FORKLAUNCH_API_TOKEN)\n\
             Set FORKLAUNCH_FORCE_DEVICE_LOGIN=1 to run the device flow anyway."
        );
    }

    // Step 1: Request device code
    log_info!(stdout, "Forklaunch CLI Login");
    log_info!(stdout, "Requesting device authorization...");

    let client = reqwest::blocking::Client::new();
    let device_response = client
        .post(format!("{}/api/auth/device/code", api_url))
        .json(&serde_json::json!({
            "client_id": "forklaunch-cli",
            "scope": "openid profile email"
        }))
        .send()?;

    if !device_response.status().is_success() {
        bail!(
            "Failed to request device code: {}",
            device_response.status()
        );
    }

    let device_data: DeviceCodeResponse = device_response.json()?;

    // Step 2: Display user code and open browser
    writeln!(stdout)?;
    log_header!(
        stdout,
        Color::Yellow,
        "Please visit: {}",
        device_data.verification_uri
    );
    log_header!(
        stdout,
        Color::Yellow,
        "Enter code: {}",
        device_data.user_code
    );
    writeln!(stdout)?;

    // Try to open browser
    let url_to_open = device_data
        .verification_uri_complete
        .as_ref()
        .unwrap_or(&device_data.verification_uri);

    log_info!(stdout, "Opening browser...");

    if let Err(e) = opener::open(url_to_open) {
        log_warn!(stdout, "Could not open browser automatically: {}", e);
        log_warn!(stdout, "Please open the URL manually.");
    }

    // Step 3: Poll for token
    let interval = Duration::from_secs(device_data.interval.unwrap_or(5) as u64);
    let mut polling_interval = interval;

    log_info!(stdout, "Waiting for authorization...");

    loop {
        sleep(polling_interval);

        let token_response = client
            .post(format!("{}/api/auth/device/token", api_url))
            .json(&serde_json::json!({
                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                "device_code": device_data.device_code,
                "client_id": "forklaunch-cli"
            }))
            .send()?;

        if token_response.status().is_success() {
            // Got session token from device auth - now exchange for JWT
            let response_body = token_response.text()?;

            let token_data: TokenResponse = serde_json::from_str(&response_body)?;
            let session_token = token_data.access_token;

            // Call /api/auth/token with session cookie to get JWT (same as browser)
            log_info!(stdout, "Exchanging session for JWT...");

            let jwt_url = format!("{}/api/auth/token", api_url);

            let jwt_response = client
                .get(&jwt_url)
                // Use Bearer auth with the session token (supported by bearer plugin)
                .bearer_auth(&session_token)
                .send()?;

            if !jwt_response.status().is_success() {
                let status = jwt_response.status();
                let body = jwt_response.text().unwrap_or_default();
                bail!("Failed to get JWT: {} - {}", status, body);
            }

            let jwt_body = jwt_response.text()?;

            #[derive(Deserialize)]
            struct JwtResponse {
                token: String,
                #[serde(rename = "expiresIn")]
                expires_in: Option<i64>,
            }

            let jwt_data: JwtResponse = serde_json::from_str(&jwt_body)?;
            let expires_at = chrono::Utc::now().timestamp() + jwt_data.expires_in.unwrap_or(604800);

            let token_storage = TokenData {
                access_token: jwt_data.token,
                refresh_token: session_token, // Use session token as refresh token
                expires_at,
            };

            let token_path = get_token_path()?;

            // Ensure parent directory exists with owner-only permissions (0o700)
            if let Some(parent) = token_path.parent() {
                #[cfg(unix)]
                {
                    use std::fs::DirBuilder;
                    let mut builder = DirBuilder::new();
                    builder.recursive(true);
                    builder.mode(0o700);
                    builder.create(parent)?;
                }

                #[cfg(not(unix))]
                {
                    create_dir_all(parent)?;
                }
            }

            let toml_content = toml::to_string(&token_storage)?;

            // Write token file with owner-only permissions (0o600)
            #[cfg(unix)]
            {
                use std::io::Write as IoWrite;

                let mut file = OpenOptions::new()
                    .write(true)
                    .create(true)
                    .truncate(true)
                    .mode(0o600)
                    .open(&token_path)?;

                file.write_all(toml_content.as_bytes())?;
            }

            #[cfg(not(unix))]
            {
                use std::fs::write;
                write(&token_path, toml_content)?;
            }

            writeln!(stdout)?;
            log_header!(stdout, Color::Green, "Successfully logged in!");

            return Ok(());
        } else {
            let error_data: Result<TokenErrorResponse, _> = token_response.json();

            match error_data {
                Ok(error) => match error.error.as_str() {
                    "authorization_pending" => {
                        continue;
                    }
                    "slow_down" => {
                        polling_interval += Duration::from_secs(5);
                        log_warn!(
                            stdout,
                            "Slowing down polling to {}s",
                            polling_interval.as_secs()
                        );
                        continue;
                    }
                    "access_denied" => {
                        bail!("Access was denied by the user");
                    }
                    "expired_token" => {
                        bail!("The device code has expired. Please try again.");
                    }
                    _ => {
                        bail!("Error: {}", error.error_description.unwrap_or(error.error));
                    }
                },
                Err(_) => {
                    bail!("Failed to authenticate: unexpected response");
                }
            }
        }
    }
}

impl CliCommand for LoginCommand {
    fn command(&self) -> Command {
        command("login", "Login to the forklaunch platform")
            .arg(
                Arg::new("token")
                    .long("token")
                    .short('t')
                    .value_name("API_TOKEN")
                    .help("API token for headless authentication (for CI/CD). Can also be set via FORKLAUNCH_API_TOKEN environment variable"),
            )
    }

    fn handler(&self, matches: &ArgMatches) -> Result<()> {
        if let Some(token) = matches.get_one::<String>("token") {
            return login_with_token(token);
        }

        if let Ok(token) = std::env::var("FORKLAUNCH_API_TOKEN") {
            return login_with_token(&token);
        }

        login()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The decision itself, without touching process-wide state: a terminal
    /// on both ends, or the explicit override.
    #[test]
    fn the_device_flow_needs_a_terminal_or_the_override() {
        assert!(device_login_decision(false, true, true));
        assert!(!device_login_decision(false, false, true));
        assert!(!device_login_decision(false, true, false));
        assert!(device_login_decision(true, false, false));
    }

    /// Under `cargo test` stdin is not a terminal — the CI shape — so `login`
    /// refuses immediately and names the headless route instead of polling
    /// for ten minutes.
    #[test]
    fn login_without_a_terminal_fails_fast_and_says_what_to_do() {
        if device_login_is_usable() {
            return; // a developer running the suite from a real terminal
        }
        let err = login().unwrap_err().to_string();
        assert!(err.contains("--token"), "{err}");
        assert!(err.contains("FORKLAUNCH_API_TOKEN"), "{err}");
    }
}