zc2 0.0.30

//! `zc up | share | down | price` while the agent runs (spec §6.1): one owner.

use crate::agent::client::AgentClient;
use crate::agent::hub::{self, HubClient, HubError};
use std::path::PathBuf;
use std::time::Duration;

const T: Duration = Duration::from_secs(5);

/// `--workers N` / `-w N` / `-n N` / `--workers=N`, only when given explicitly.
pub fn explicit_workers(args: &[String]) -> Option<u32> {
    let mut it = args.iter();
    while let Some(a) = it.next() {
        match a.as_str() {
            "--workers" | "-w" | "-n" => return it.next().and_then(|v| v.parse().ok()),
            s if s.starts_with("--workers=") => return s["--workers=".len()..].parse().ok(),
            _ => {}
        }
    }
    None
}

pub fn port_flags_given(args: &[String]) -> bool {
    args.iter()
        .any(|a| matches!(a.as_str(), "--port" | "-p" | "--broker-port" | "-b"))
}

fn num(v: f64) -> String {
    // 18.0 → "18", 3.6 → "3.6"
    format!("{v}")
}

pub fn format_price(state: &str, per_hour: Option<f64>, effective: Option<f64>) -> String {
    match (state, per_hour, effective) {
        ("set", Some(p), _) => format!("this Mac: {} credits/hour", num(p)),
        ("inherited", _, Some(e)) => format!("this Mac: account default ({} credits/hour)", num(e)),
        ("inherited", _, None) => "this Mac: account default".to_string(),
        ("mixed", _, Some(e)) => format!("this Mac: mixed (from {} credits/hour)", num(e)),
        ("disabled", _, _) => "this Mac: disabled on the hub".to_string(),
        _ => "this Mac: no price yet".to_string(),
    }
}

fn print_error(e: String) -> i32 {
    eprintln!("zc: {e}");
    1
}

pub fn up(c: &AgentClient, args: &[String]) -> i32 {
    if port_flags_given(args) {
        eprintln!("  note: the zc agent owns the ports; --port/--broker-port are ignored");
    }
    if let Some(n) = explicit_workers(args) {
        if let Err(e) = c.put("/v1/workers", serde_json::json!({ "count": n }), T) {
            return print_error(e);
        }
    }
    match c.put("/v1/sharing", serde_json::json!({ "on": true }), T) {
        Ok(s) => {
            println!(
                "✓ Sharing handed to zc agent: {} worker(s). Follow it with `zc agent status`.",
                s["this_mac"]["workers"]["desired"]
            );
            0
        }
        Err(e) => print_error(e),
    }
}

/// `zc down` while the agent runs: turn sharing off through the agent. A
/// broker the agent didn't start (a manual `zc up -d`) is one it can't stop,
/// so then this returns `None` once sharing is off: the caller falls through
/// to the port-based teardown that stops that broker, which is what the
/// `unmanaged_broker` hint and the 409 tell the owner to run `zc down` for.
pub fn down(c: &AgentClient) -> Option<i32> {
    match c.put("/v1/sharing", serde_json::json!({ "on": false }), T) {
        Ok(s) if s["this_mac"]["broker"] == "unmanaged" => {
            println!("✓ zc agent paused sharing; stopping the broker it didn't start.");
            None
        }
        Ok(_) => {
            println!("✓ zc agent is draining the workers and will stop sharing.");
            Some(0)
        }
        Err(e) => Some(print_error(e)),
    }
}

/// `nan` and `inf` parse as numbers, but JSON has no such values: serialized
/// they become `null`, which the hub reads as "use the account default".
fn finite_price(v: f64) -> Result<f64, String> {
    if v.is_finite() {
        Ok(v)
    } else {
        Err(format!(
            "price must be a finite number of credits/hour, got {v}"
        ))
    }
}

pub fn price_set(c: &AgentClient, value: f64) -> i32 {
    if let Err(e) = finite_price(value) {
        return print_error(e);
    }
    let body = serde_json::json!({ "scope": "device", "price_per_hour": value });
    match c.put("/v1/price", body, Duration::from_secs(30)) {
        Ok(_) => {
            println!("price set: {} credits/hour on this Mac (hub)", num(value));
            0
        }
        Err(e) => print_error(e),
    }
}

/// `zc price inherit` — drop this Mac's own price so it follows the account
/// default again.
pub fn price_clear(c: &AgentClient) -> i32 {
    // JSON null is what tells the hub "no price of its own"; a number would
    // pin this Mac to whatever it happened to be.
    let body = serde_json::json!({ "scope": "device", "price_per_hour": null });
    match c.put("/v1/price", body, Duration::from_secs(30)) {
        Ok(_) => {
            println!("price cleared: this Mac follows the account default (hub)");
            0
        }
        Err(e) => print_error(e),
    }
}

/// `zc price default <v>` — the account default, which every device without
/// its own price follows.
pub fn price_default(c: &AgentClient, value: f64) -> i32 {
    if let Err(e) = finite_price(value) {
        return print_error(e);
    }
    let body = serde_json::json!({ "scope": "default", "price_per_hour": value });
    match c.put("/v1/price", body, Duration::from_secs(30)) {
        Ok(_) => {
            println!(
                "account default set: {} credits/hour, on every device with no price of its own (hub)",
                num(value)
            );
            0
        }
        Err(e) => print_error(e),
    }
}

pub fn price_show(c: &AgentClient) -> i32 {
    match c.summary(T) {
        Ok(s) => {
            let p = &s["prices"]["this_mac"];
            println!(
                "{}",
                format_price(
                    p["state"].as_str().unwrap_or(""),
                    p["per_hour"].as_f64(),
                    p["effective_per_hour"].as_f64()
                )
            );
            0
        }
        Err(e) => print_error(e),
    }
}

pub fn try_up(args: &[String]) -> Option<i32> {
    AgentClient::running().map(|c| up(&c, args))
}

pub fn try_down() -> Option<i32> {
    AgentClient::running().and_then(|c| down(&c))
}

pub fn try_price(cmd: crate::PriceCmd) -> Option<i32> {
    use crate::PriceCmd::*;
    let c = AgentClient::running()?;
    Some(match cmd {
        Show => price_show(&c),
        Device(v) => price_set(&c, v),
        Inherit => price_clear(&c),
        Default(v) => price_default(&c, v),
    })
}

/// `GET /api/provide/summary`, retried once in UTC when the hub rejects the
/// starting time zone with 422 (mirrors the agent's own `refresh_hub` retry).
async fn summary_with_retry(
    client: &HubClient,
    tz: &str,
) -> Result<crate::agent::hubwire::HubSummary, HubError> {
    let result = client.summary(tz).await;
    if tz != "UTC" && matches!(result, Err(HubError::Status(422, _))) {
        return client.summary("UTC").await;
    }
    result
}

/// `zc price [<v>]` without an agent: the same hub writes it would make.
pub fn price_direct(cmd: crate::PriceCmd) -> i32 {
    price_direct_in(crate::credentials::dir(), cmd)
}

/// `price_direct`, with its state dir (`~/.zakuro`: the credentials and the
/// node key) passed in.
fn price_direct_in(state_dir: Option<PathBuf>, cmd: crate::PriceCmd) -> i32 {
    // Before anything is read or sent.
    if let Some(Err(e)) = cmd.value().map(finite_price) {
        return print_error(e);
    }

    let Some(state_dir) = state_dir else {
        return print_error("no HOME or ZAKURO_HOME".into());
    };
    let Some(creds) = hub::read_credentials(&state_dir.join("credentials")) else {
        return print_error("this Mac is not signed in: run `zc login`".into());
    };
    let base = creds
        .api_url
        .clone()
        .unwrap_or_else(crate::credentials::default_api_url);
    let pk = crate::broker::node_identity::NodeKey::load_or_create_in(Some(state_dir)).public_b64();
    let rt = match tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
    {
        Ok(rt) => rt,
        Err(e) => return print_error(e.to_string()),
    };
    rt.block_on(async move {
        let client = HubClient::new(&base, &creds.api_key, Duration::from_secs(10));
        let summary = match summary_with_retry(&client, &hub::system_tz()).await {
            Ok(s) => s,
            Err(e) => return print_error(e.message()),
        };
        let mine = summary.this_mac_workers(&pk);
        let b = summary.bounds();
        if let Some(v) = cmd.value() {
            if v < b.min as f64 || v > b.max as f64 {
                return print_error(format!(
                    "price must be between {} and {} credits/hour",
                    b.min, b.max
                ));
            }
        }

        // The account default is account-wide: no worker row is involved.
        if let crate::PriceCmd::Default(v) = cmd {
            return match client.put_default_price(Some(v)).await {
                Ok(()) => {
                    println!(
                        "account default set: {} credits/hour, on every device with no price of its own",
                        num(v)
                    );
                    0
                }
                Err(e) => print_error(e.message()),
            };
        }

        // What to write on this Mac's worker rows: a price, or null to follow
        // the account default.
        let price = match cmd {
            crate::PriceCmd::Show => {
                let p = crate::agent::merge::this_mac_price(&mine);
                println!(
                    "{}",
                    format_price(&p.state, p.per_hour, p.effective_per_hour)
                );
                return 0;
            }
            crate::PriceCmd::Device(v) => Some(v),
            crate::PriceCmd::Inherit => None,
            crate::PriceCmd::Default(_) => unreachable!("returned just above"),
        };

        let ids: Vec<i64> = mine.iter().map(|w| w.id).collect();
        if ids.is_empty() {
            return print_error(
                "this Mac has no workers on the hub yet: start sharing first (`zc share`)".into(),
            );
        }
        match hub::set_device_price(&client, &ids, price).await {
            Ok(()) => {
                match price {
                    Some(v) => println!(
                        "price set: {} credits/hour on this Mac ({} worker(s))",
                        num(v),
                        ids.len()
                    ),
                    None => println!(
                        "price cleared: this Mac follows the account default ({} worker(s))",
                        ids.len()
                    ),
                }
                0
            }
            Err(failed) => print_error(format!("could not price {} worker(s)", failed.len())),
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::client::tests::{idle_agent, sleeper_cfg, start_agent};
    use crate::agent::core::tests::{hub_json, row};
    use crate::agent::hub::tests::{mock_hub, Seen};

    fn args(s: &[&str]) -> Vec<String> {
        s.iter().map(|a| a.to_string()).collect()
    }

    #[test]
    fn worker_and_port_flags_are_detected() {
        assert_eq!(explicit_workers(&args(&["--workers", "3"])), Some(3));
        assert_eq!(explicit_workers(&args(&["-w", "2", "-d"])), Some(2));
        assert_eq!(explicit_workers(&args(&["-n", "4"])), Some(4));
        assert_eq!(explicit_workers(&args(&["--workers=5"])), Some(5));
        assert_eq!(explicit_workers(&args(&["-d"])), None);
        assert!(port_flags_given(&args(&["--port", "4000"])));
        assert!(port_flags_given(&args(&["-b", "9100"])));
        assert!(!port_flags_given(&args(&["--workers", "2"])));
    }

    #[test]
    fn price_lines() {
        assert_eq!(
            format_price("set", Some(18.0), Some(18.0)),
            "this Mac: 18 credits/hour"
        );
        assert_eq!(
            format_price("inherited", None, Some(12.0)),
            "this Mac: account default (12 credits/hour)"
        );
        assert_eq!(
            format_price("inherited", None, None),
            "this Mac: account default"
        );
        assert_eq!(
            format_price("mixed", None, Some(12.0)),
            "this Mac: mixed (from 12 credits/hour)"
        );
    }

    #[test]
    fn a_disabled_price_state_is_shown_as_disabled_on_the_hub() {
        assert_eq!(
            format_price("disabled", None, None),
            "this Mac: disabled on the hub"
        );
    }

    #[test]
    fn a_422_on_the_starting_zone_is_retried_once_in_utc() {
        use crate::agent::hubwire::tests::HUB_JSON;

        let (base, seen) = mock_hub(|_, url, _| match url {
            "/api/provide/summary?tz=Asia%2FTokyo" => (422, "{}".into()),
            "/api/provide/summary?tz=UTC" => (200, HUB_JSON.to_string()),
            _ => (404, "{}".into()),
        });
        let client = hub::HubClient::new(&base, "zk_1_x", std::time::Duration::from_secs(5));
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let summary = rt
            .block_on(summary_with_retry(&client, "Asia/Tokyo"))
            .expect("retried in UTC and succeeded");
        assert_eq!(summary.account.username, "jean");
        let urls: Vec<String> = seen
            .lock()
            .unwrap()
            .iter()
            .map(|(_, u, _, _)| u.clone())
            .collect();
        assert_eq!(
            urls,
            vec![
                "/api/provide/summary?tz=Asia%2FTokyo".to_string(),
                "/api/provide/summary?tz=UTC".to_string()
            ]
        );
    }

    #[test]
    fn up_and_down_are_handed_to_the_agent() {
        let (c, _stop) = idle_agent();
        assert_eq!(up(&c, &args(&["--workers", "2", "--port", "4000"])), 0);
        let s = c.summary(std::time::Duration::from_secs(2)).unwrap();
        assert_eq!(
            (
                s["this_mac"]["sharing"].as_bool(),
                s["this_mac"]["workers"]["desired"].as_u64()
            ),
            (Some(true), Some(2))
        );
        assert_eq!(down(&c), Some(0));
        assert_eq!(
            c.summary(std::time::Duration::from_secs(2)).unwrap()["this_mac"]["sharing"],
            false
        );
        assert_eq!(
            price_set(&c, 10.0),
            1,
            "not signed in: the agent answers 503"
        );
    }

    fn wait_until(what: &str, mut done: impl FnMut() -> bool) {
        let end = std::time::Instant::now() + Duration::from_secs(10);
        while std::time::Instant::now() < end {
            if done() {
                return;
            }
            std::thread::sleep(Duration::from_millis(100));
        }
        panic!("timed out waiting for: {what}");
    }

    /// A broker the agent didn't start (a manual `zc up -d`) holds the broker
    /// port. `zc down` pauses the agent, then hands back to the port-based
    /// teardown that stops that broker, so following the `unmanaged_broker`
    /// hint (and the 409) actually frees this Mac for the agent.
    #[test]
    fn down_with_an_unmanaged_broker_pauses_the_agent_then_falls_through() {
        // Answers `/health` like a zakuro broker, but the agent never started it.
        let (fake_broker, _) = mock_hub(|_, url, _| match url {
            "/health" => (
                200,
                r#"{"status":"healthy","service":"zakuro-broker"}"#.into(),
            ),
            _ => (404, "{}".into()),
        });
        let mut cfg = sleeper_cfg("down-unmanaged");
        cfg.broker_port = fake_broker.rsplit(':').next().unwrap().parse().unwrap();
        crate::agent::files::ensure_dir(&cfg.dir).unwrap();
        crate::agent::files::save_json(
            &cfg.dir.join(crate::agent::files::STATE_FILE),
            &crate::agent::files::DesiredState {
                sharing: true,
                ..Default::default()
            },
        )
        .unwrap();
        let (c, _guard) = start_agent(cfg);
        let summary = || c.summary(Duration::from_secs(2)).unwrap();
        wait_until("the agent sees the unmanaged broker", || {
            summary()["this_mac"]["broker"] == "unmanaged"
        });
        assert_eq!(summary()["this_mac"]["sharing"], true);

        assert_eq!(down(&c), None, "the port-based `zc down` must still run");
        assert_eq!(
            summary()["this_mac"]["sharing"],
            false,
            "the agent paused sharing first"
        );
    }

    /// A hub whose summary lists this node's one worker row (11); every PUT
    /// succeeds.
    fn pricing_hub(pk: String) -> (String, Seen) {
        mock_hub(move |method, url, _| match method {
            "GET" if url.starts_with("/api/provide/summary") => {
                (200, hub_json(&pk, &row(11, None)))
            }
            "PUT" => (200, "{}".into()),
            _ => (404, "{}".into()),
        })
    }

    /// (url, body) of every PUT: the account default and a device price go to
    /// different routes, so a test must see which one was written.
    fn put_calls(seen: &Seen) -> Vec<(String, String)> {
        seen.lock()
            .unwrap()
            .iter()
            .filter(|e| e.0 == "PUT")
            .map(|e| (e.1.clone(), e.2.clone()))
            .collect()
    }

    /// A signed-in state dir pointed at `hub`, with this Mac's node key.
    fn signed_in(name: &str) -> (PathBuf, String, Seen) {
        let state = crate::agent::files::tests::tmp(name);
        std::fs::create_dir_all(&state).unwrap();
        let pk = crate::broker::node_identity::NodeKey::load_or_create_in(Some(state.clone()))
            .public_b64();
        let (hub, seen) = pricing_hub(pk);
        std::fs::write(
            state.join("credentials"),
            format!("api_key=zk_1_test\napi_url={hub}\n"),
        )
        .unwrap();
        (state, hub, seen)
    }

    /// `zc price inherit` must send JSON null on this Mac's worker rows: that
    /// is what tells the hub to fall back to the account default. Sending a
    /// number, or nothing, would leave the Mac pinned to its old price.
    #[test]
    fn inherit_clears_this_macs_own_price() {
        let (state, _hub, seen) = signed_in("price-inherit");
        assert_eq!(price_direct_in(Some(state), crate::PriceCmd::Inherit), 0);
        assert_eq!(
            put_calls(&seen),
            vec![(
                "/api/workers/11/price".to_string(),
                r#"{"price_per_hour":null}"#.to_string()
            )]
        );
    }

    /// `zc price default <v>` is an account-wide write, so it goes to the
    /// global-price route and must not touch this Mac's worker rows.
    #[test]
    fn the_account_default_is_written_account_wide() {
        let (state, _hub, seen) = signed_in("price-default");
        assert_eq!(
            price_direct_in(Some(state), crate::PriceCmd::Default(42.0)),
            0
        );
        assert_eq!(
            put_calls(&seen),
            vec![(
                "/api/workers/global-price".to_string(),
                r#"{"default_price_per_hour":42.0}"#.to_string()
            )]
        );
    }

    /// The same two commands through a running agent: same hub writes.
    #[test]
    fn the_agent_makes_the_same_two_writes() {
        let cfg = sleeper_cfg("price-agent-cmds");
        std::fs::create_dir_all(&cfg.state_dir).unwrap();
        let (hub, seen) = pricing_hub(cfg.node_pubkey());
        std::fs::write(
            cfg.state_dir.join("credentials"),
            format!("api_key=zk_1_test\napi_url={hub}\n"),
        )
        .unwrap();
        let (c, _guard) = start_agent(cfg);

        assert_eq!(price_clear(&c), 0, "inherit");
        assert_eq!(price_default(&c, 42.0), 0, "default");
        assert_eq!(
            put_calls(&seen),
            vec![
                (
                    "/api/workers/11/price".to_string(),
                    r#"{"price_per_hour":null}"#.to_string()
                ),
                (
                    "/api/workers/global-price".to_string(),
                    r#"{"default_price_per_hour":42.0}"#.to_string()
                ),
            ]
        );
    }

    fn put_bodies(seen: &Seen) -> Vec<String> {
        seen.lock()
            .unwrap()
            .iter()
            .filter(|e| e.0 == "PUT")
            .map(|e| e.2.clone())
            .collect()
    }

    /// `nan` and `inf` serialize as JSON `null`, which the hub reads as "use
    /// the account default": `zc price nan` would silently clear this Mac's
    /// price. Without an agent it is refused before anything is sent.
    #[test]
    fn a_non_finite_price_is_refused_before_it_reaches_the_hub() {
        let state = crate::agent::files::tests::tmp("price-direct");
        std::fs::create_dir_all(&state).unwrap();
        let pk = crate::broker::node_identity::NodeKey::load_or_create_in(Some(state.clone()))
            .public_b64();
        let (hub, seen) = pricing_hub(pk);
        std::fs::write(
            state.join("credentials"),
            format!("api_key=zk_1_test\napi_url={hub}\n"),
        )
        .unwrap();

        for v in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
            assert_eq!(
                price_direct_in(Some(state.clone()), crate::PriceCmd::Device(v)),
                1,
                "{v}"
            );
        }
        assert_eq!(put_bodies(&seen), Vec::<String>::new(), "nothing sent");

        assert_eq!(
            price_direct_in(Some(state), crate::PriceCmd::Device(18.0)),
            0,
            "a real price still goes through"
        );
        assert_eq!(
            put_bodies(&seen),
            vec![r#"{"price_per_hour":18.0}"#.to_string()]
        );
    }

    /// The same through a running agent: refused before the request is sent.
    #[test]
    fn a_non_finite_price_is_refused_before_it_reaches_the_agent() {
        let cfg = sleeper_cfg("price-set");
        std::fs::create_dir_all(&cfg.state_dir).unwrap();
        let (hub, seen) = pricing_hub(cfg.node_pubkey());
        std::fs::write(
            cfg.state_dir.join("credentials"),
            format!("api_key=zk_1_test\napi_url={hub}\n"),
        )
        .unwrap();
        let (c, _guard) = start_agent(cfg);

        for v in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
            assert_eq!(price_set(&c, v), 1, "{v}");
        }
        assert_eq!(put_bodies(&seen), Vec::<String>::new(), "nothing sent");

        assert_eq!(price_set(&c, 18.0), 0, "a real price still goes through");
        assert_eq!(
            put_bodies(&seen),
            vec![r#"{"price_per_hour":18.0}"#.to_string()]
        );
    }
}