zc2 0.0.29

//! Broker module for routing compute requests to workers.
//!
//! The broker provides:
//! - Worker registry with health monitoring
//! - Credit-based billing system with dashboard API
//! - Optimal worker selection based on price-per-compute
//! - Low-latency request forwarding
//! - WireGuard-based P2P worker discovery
//! - Live terminal UI for monitoring
//! - Benchmarking tools

pub mod bench;
pub mod bench_mesh;
pub mod credits;
pub mod deploy;
pub mod discovery;
pub mod flush;
pub mod http_adapter;
pub mod info;
pub mod ledger;
pub mod node_health;
pub mod node_identity;
pub mod node_sync;
pub mod peer;
pub mod pricing;
pub mod quic;
pub mod recovery;
pub mod roster_cache;
pub mod router;
pub mod selection;
pub mod server;
pub mod stats;
pub mod task_board;
pub mod uri;
pub mod voucher;
pub mod wal;
pub mod worker;
pub mod worker_quic;

use std::sync::Arc;
use std::sync::OnceLock;

use dashmap::DashMap;

pub use credits::CreditManager;
pub use discovery::DiscoveryConfig;
pub use flush::TransactionBuffer;
pub use ledger::Ledger;
pub use peer::PeerManager;
pub use pricing::BrokerPrice;
pub use router::Router;
pub use server::start_server;
pub use stats::StatsCollector;
pub use wal::Wal;
pub use worker::WorkerRegistry;

/// Infer this node's name from the system hostname.
///
/// This used to query the mesh daemon first -- `wireguard status --json`, then
/// its local API socket over curl -- and fall back to the hostname. WireGuard
/// has no such daemon and no name service: a peer is an address, and the node
/// name is the host's own. So the fallback is now the whole implementation,
/// and two subprocess spawns per call disappear with it.
fn detect_node_name() -> Option<String> {
    std::process::Command::new("hostname")
        .output()
        .ok()
        .and_then(|o| String::from_utf8(o.stdout).ok())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
}

/// This broker's node name (ZAKURO_NODE_NAME or detected), or "unknown".
pub fn node_name_or_default() -> String {
    std::env::var("ZAKURO_NODE_NAME")
        .ok()
        .filter(|s| !s.is_empty())
        .or_else(detect_node_name)
        .unwrap_or_else(|| "unknown".to_string())
}

/// Broker configuration
#[derive(Debug, Clone)]
pub struct BrokerConfig {
    /// Host to bind the broker server
    pub host: String,
    /// Port to bind the broker server
    pub port: u16,
    /// Worker health check interval in seconds
    pub health_check_interval: u64,
    /// Worker timeout before marking unhealthy (seconds)
    pub worker_timeout: u64,
    /// Minimum credits required for any request
    pub min_credits: f64,
    /// Run in daemon mode (background, minimal output)
    pub daemon: bool,
    /// Verbose output (show live transactions)
    pub verbose: bool,
    /// Enable WireGuard-based worker discovery
    pub enable_discovery: bool,
    /// Discovery configuration
    pub discovery: DiscoveryConfig,
    /// Shared secret for worker management endpoints (from ZAKURO_WORKER_KEY env var)
    pub worker_key: Option<String>,
    /// Owner zakuro_user_id — workers discovered by this broker belong to this user
    /// (from ZAKURO_OWNER_ID env var, required for PG worker sync)
    pub owner_user_id: Option<String>,
    /// Human-readable node name for this broker instance (from ZAKURO_NODE_NAME env var)
    /// Used to tag synced workers with their source node.
    pub node_name: Option<String>,
    /// Shared secret for peer-to-peer broker communication (from ZAKURO_PEER_KEY env var)
    pub peer_key: Option<String>,
    /// Enable P2P credit operations (from ZAKURO_P2P env var, default false)
    pub enable_p2p: bool,
    /// Dashboard API URL for worker sync (from ZAKURO_API_URL env var)
    pub api_url: Option<String>,
    /// User API key for authentication (from ZAKURO_API_KEY env var)
    /// This is the universal key used for both user services and broker-to-dashboard communication
    pub api_key: Option<String>,
    /// Override WireGuard IP detection (test-only, avoids global env var race).
    /// When set, `BrokerState::own_wireguard_ip` uses this value instead of
    /// probing network interfaces.
    pub wireguard_ip_override: Option<String>,
    /// Explicit QUIC port (UDP).  When `None`, defaults to `port + 1`.
    pub quic_port: Option<u16>,
    /// Override the async runtime's worker-thread count (test-only).
    /// `None` → tokio default (num_cpus). Tests set a small value so the many
    /// broker runtimes a parallel `cargo test` leaks don't oversubscribe the
    /// CPU and flake latency-sensitive tests. See `server::build_broker_runtime`.
    pub runtime_worker_threads: Option<usize>,
}

impl Default for BrokerConfig {
    fn default() -> Self {
        Self {
            host: "0.0.0.0".to_string(),
            port: 9000,
            health_check_interval: 5,
            worker_timeout: 30,
            min_credits: 0.001,
            daemon: false,
            verbose: true,
            enable_discovery: true,
            discovery: DiscoveryConfig::default(),
            worker_key: std::env::var("ZAKURO_WORKER_KEY").ok(),
            owner_user_id: {
                // Always derived from API key format: zk_{user_id}_{hex}.
                // Not user-configurable — ZAKURO_OWNER_ID is intentionally ignored.
                std::env::var("ZAKURO_API_KEY").ok().and_then(|k| {
                    k.strip_prefix("zk_")
                        .and_then(|rest| rest.rfind('_').map(|pos| rest[..pos].to_string()))
                })
            },
            node_name: std::env::var("ZAKURO_NODE_NAME")
                .ok()
                .or_else(detect_node_name),
            // Env first; else the key `zc connect` stored from the hub, so a
            // user's broker peers with the mesh without pasting the secret.
            peer_key: std::env::var("ZAKURO_PEER_KEY")
                .ok()
                .filter(|k| !k.trim().is_empty())
                .or_else(crate::credentials::load_mesh_peer_key),
            enable_p2p: std::env::var("ZAKURO_P2P")
                .map(|v| v == "true" || v == "1")
                .unwrap_or(false),
            api_url: std::env::var("ZAKURO_API_URL").ok(),
            api_key: std::env::var("ZAKURO_API_KEY").ok(),
            wireguard_ip_override: None,
            quic_port: None,
            runtime_worker_threads: None,
        }
    }
}

/// P2P for a broker started from the CLI (`zc broker`, `zc share`): honour an
/// explicit `ZAKURO_P2P`, otherwise turn it on exactly when this machine is on
/// the mesh — a broker with a `10.13.13.x` address that does not peer would be
/// reachable by everyone and dial nobody. `BrokerConfig::default()` itself
/// stays env-only so tests constructing configs are not swayed by the host.
pub fn p2p_default() -> bool {
    match std::env::var("ZAKURO_P2P") {
        Ok(v) => v == "true" || v == "1",
        Err(_) => discovery::get_mesh_ip().is_some(),
    }
}

/// Defaults for a broker started from the CLI on a user's machine (`zc broker`,
/// `zc share`, the auto-spawned local broker). Library code reads plain env
/// vars, so these are applied to the process environment BEFORE the config is
/// built or a child is spawned (children inherit them):
///
/// - `ZAKURO_REQUIRE_VOUCHER=1` whenever billing credentials are present and
///   the operator has not set it. The fleet runs with the flag on and rejects
///   unvouched offers with 402, while a requesting broker only mints when ITS
///   flag is on — so a user broker with the default left unset could never
///   dispatch to the fleet, and the failure looked like a dead mesh.
pub fn apply_user_broker_defaults() {
    crate::credentials::load_into_env();
    let set = |k: &str| {
        std::env::var(k)
            .map(|v| !v.trim().is_empty())
            .unwrap_or(false)
    };
    if set("ZAKURO_API_KEY")
        && set("ZAKURO_API_URL")
        && std::env::var("ZAKURO_REQUIRE_VOUCHER").is_err()
    {
        std::env::set_var("ZAKURO_REQUIRE_VOUCHER", "1");
    }
}

/// Main broker state shared across handlers
pub struct BrokerState {
    /// Worker registry
    pub workers: WorkerRegistry,
    /// Credit manager (legacy, kept for compatibility)
    pub credits: CreditManager,
    /// Central ledger for credit management
    pub ledger: Ledger,
    /// Router for worker selection
    pub router: Router,
    /// Configuration
    pub config: BrokerConfig,
    /// Active request tracking (request_id -> worker_id)
    pub active_requests: DashMap<String, String>,
    /// Instance affinity registry (instance_id -> worker_id)
    /// Used to pin RemoteProxy instances to the worker that created them.
    pub instance_registry: DashMap<String, String>,
    /// Whether running in local mode (free execution)
    pub local_mode: std::sync::atomic::AtomicBool,
    /// Statistics collector behind `/stats`
    pub stats: Arc<StatsCollector>,
    /// Write-ahead log for crash recovery
    pub wal: Wal,
    /// This node's own WireGuard IP (for detecting local vs remote workers)
    pub own_wireguard_ip: Option<String>,
    /// P2P peer manager for broker-to-broker credit operations
    pub peer_manager: PeerManager,
    /// Transaction flush buffer (batched PG writes in P2P mode)
    pub tx_buffer: TransactionBuffer,
    /// QUIC transport for high-throughput peer task offers (None if not started)
    pub quic: OnceLock<Arc<quic::QuicTransport>>,
    /// Actual QUIC port this broker is listening on (set after QUIC starts)
    pub quic_port: std::sync::atomic::AtomicU16,
    /// Shared connection-pooled HTTP client for forwarding to workers. Reused
    /// across requests so each /execute no longer pays a fresh TCP+TLS setup
    /// (audit M3/async-phase-2: connection pooling). Per-request timeouts are
    /// set on the request builder, not the agent.
    pub http_client: ureq::Agent,
    /// Shared, connection-pooled async client for the /execute worker forward.
    pub async_http: reqwest::Client,
    /// Two-phase peer offers reserved but not yet committed/cancelled, keyed by
    /// task_id. A reserve holds a worker quota slot; commit executes, cancel (or
    /// TTL sweep) releases. See server::{reserve,cancel,take_committed,sweep}.
    pub pending_offers: DashMap<String, PendingOffer>,
    /// This node's Ed25519 signing identity — signs every outbound /peer/*
    /// request and is the zero-trust replacement for the shared peer key.
    pub node_key: Arc<node_identity::NodeKey>,
    /// Dashboard-synced roster of authorized node pubkeys → `revoked` flag.
    /// Empty = permissive (no signed peers known yet); populated by roster sync.
    pub node_roster: DashMap<String, bool>,
    /// Anti-replay guard for inbound signed /peer/* requests.
    pub node_sig_guard: node_identity::ReplayGuard,
    /// Cached dashboard voucher-signing pubkey (b64) for offline voucher verify.
    pub dash_voucher_pubkey: std::sync::RwLock<Option<String>>,
    /// This broker's advertised price `P` (credits/hour). Live-mutable via
    /// `zc price <value>` / the local-only `POST /price` control endpoint.
    pub price: BrokerPrice,
}

/// A peer task offer that was RESERVED (worker slot held) but not yet executed,
/// awaiting a two-phase commit or cancel.
pub struct PendingOffer {
    pub offer: task_board::TaskOffer,
    pub worker_id: String,
    pub created: std::time::Instant,
}

impl BrokerState {
    /// Create a new broker state with default configuration
    pub fn new() -> Self {
        Self::with_config(BrokerConfig::default())
    }

    /// Create a new broker state with custom configuration
    pub fn with_config(config: BrokerConfig) -> Self {
        let ledger = Ledger::new(config.api_url.clone(), config.api_key.clone());
        let wal = {
            // Candidate paths: ZAKURO_WAL_PATH env → $HOME/.zakuro/wal.jsonl → /tmp/zakuro-wal.jsonl
            let mut candidates: Vec<std::path::PathBuf> = vec![];
            if let Ok(p) = std::env::var("ZAKURO_WAL_PATH") {
                candidates.push(std::path::PathBuf::from(p));
            }
            if let Some(home) = std::env::var("HOME")
                .ok()
                .or_else(|| std::env::var("USERPROFILE").ok())
            {
                let dir = std::path::PathBuf::from(home).join(".zakuro");
                let _ = std::fs::create_dir_all(&dir);
                candidates.push(dir.join("wal.jsonl"));
            }
            candidates.push(std::path::PathBuf::from("/tmp/zakuro-wal.jsonl"));

            candidates
                .into_iter()
                .find_map(|p| Wal::open(p.to_str().unwrap_or("/tmp/zakuro-wal.jsonl")).ok())
                .expect("Failed to open WAL on any candidate path")
        };
        let own_wireguard_ip = config
            .wireguard_ip_override
            .clone()
            .or_else(discovery::get_effective_node_ip);

        // Peer broker list: from ZAKURO_PEERS, or discover on localhost when empty.
        // When using discovered brokers, ZAKURO_API_KEY is optional (P2P only, no dashboard).
        let peer_addresses: Vec<String> = if config.discovery.peers.is_empty()
            && config.enable_p2p
            && std::env::var("ZAKURO_DISCOVER_BROKER_PEERS").unwrap_or_else(|_| "true".into())
                != "false"
        {
            let mut discovered =
                discovery::discover_broker_peers_on_localhost(config.port, 9000, 9010);
            if !discovered.is_empty() {
                eprintln!(
                    "  [P2P] Discovered {} broker(s) on localhost (ZAKURO_API_KEY optional)",
                    discovered.len()
                );
            }

            // Best-effort: also probe the WireGuard mesh subnet for peer
            // brokers when this node is actually on the mesh. Strictly
            // gated behind P2P + a present mesh IP so this is a no-op
            // off-mesh; any failure here just means fewer peers, never a
            // hard error (see discover_broker_peers_on_mesh_subnet doc).
            if let Some(mesh_ip) = discovery::get_mesh_ip() {
                let mesh_peers = discovery::discover_broker_peers_on_mesh_subnet(
                    &mesh_ip,
                    config.peer_key.as_deref().unwrap_or(""),
                );
                if !mesh_peers.is_empty() {
                    eprintln!(
                        "  [P2P] Discovered {} broker(s) on the mesh subnet",
                        mesh_peers.len()
                    );
                }
                discovered.extend(mesh_peers);
            }

            discovered
        } else {
            config.discovery.peers.clone()
        };

        // This node's signing identity — shared with the peer manager so every
        // outbound /peer/* request is signed, and kept on state for inbound verify.
        let node_key = Arc::new(node_identity::NodeKey::load_or_create());

        // Initialize P2P peer manager
        let peer_key = config.peer_key.clone().unwrap_or_default();
        let peer_manager = PeerManager::new_with_node_key(
            own_wireguard_ip.as_deref(),
            &peer_addresses,
            config.port,
            peer_key,
            config.enable_p2p,
            Some(node_key.clone()),
        );

        Self {
            workers: WorkerRegistry::new(),
            credits: CreditManager::new(),
            ledger,
            router: Router::new(),
            config,
            active_requests: DashMap::new(),
            instance_registry: DashMap::new(),
            local_mode: std::sync::atomic::AtomicBool::new(false),
            stats: Arc::new(StatsCollector::new()),
            wal,
            own_wireguard_ip,
            peer_manager,
            tx_buffer: TransactionBuffer::new(),
            quic: OnceLock::new(),
            quic_port: std::sync::atomic::AtomicU16::new(0),
            // Pooled client (no global timeout — set per-request in forward_to_worker).
            http_client: ureq::Agent::new_with_config(ureq::Agent::config_builder().build()),
            async_http: reqwest::Client::builder()
                .pool_max_idle_per_host(32)
                .build()
                .expect("reqwest client builds with rustls"),
            pending_offers: DashMap::new(),
            node_key,
            node_roster: DashMap::new(),
            node_sig_guard: node_identity::ReplayGuard::new(),
            dash_voucher_pubkey: std::sync::RwLock::new(None),
            price: BrokerPrice::from_env_or_default(),
        }
    }

    /// Check if running in local mode (free execution)
    pub fn is_local_mode(&self) -> bool {
        self.local_mode.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Set local mode
    pub fn set_local_mode(&self, local: bool) {
        self.local_mode
            .store(local, std::sync::atomic::Ordering::Relaxed);
    }

    /// Billing is only meaningful when a centralized authority (dashboard API)
    /// is configured.  Without it, credits are unbacked local numbers — so we
    /// skip all reserve / commit / balance-check logic and let every execution
    /// run for free.
    pub fn is_billing_enabled(&self) -> bool {
        // API mode: dashboard URL + key configured
        (self.config.api_url.is_some() && self.config.api_key.is_some())
        // Standalone mode: ZAKURO_MASTER_KEY is set (local credit ledger)
        || !std::env::var("ZAKURO_MASTER_KEY").unwrap_or_default().is_empty()
    }

    /// Check if a worker URI belongs to this node (local = free execution)
    pub fn is_local_worker(&self, worker_uri: &str) -> bool {
        if self.is_local_mode() {
            return true; // Everything is local in local mode
        }
        // Localhost workers are always local
        if worker_uri.contains("127.0.0.1") || worker_uri.contains("localhost") {
            return true;
        }
        if let Some(ref own_ip) = self.own_wireguard_ip {
            worker_uri.contains(own_ip)
        } else {
            false
        }
    }

    /// Verify this broker's identity with the dashboard at startup.
    ///
    /// - Calls `GET /api/auth/me/api-key` to resolve the real `zakuro_user_id`.
    /// - If `owner_user_id` is not set, auto-fills it from the dashboard.
    /// - If `owner_user_id` IS set but doesn't match the API key's user, overrides
    ///   it with the real value and warns.
    /// - Returns the verified `zakuro_user_id`, or `None` if verification failed.
    pub fn verify_owner_with_dashboard(&mut self) -> Option<String> {
        let api_url = self.config.api_url.as_ref()?;
        let api_key = self.config.api_key.as_ref()?;

        let url = format!("{}/api/auth/me/api-key", api_url.trim_end_matches('/'));
        let agent = ureq::Agent::new_with_config(
            ureq::Agent::config_builder()
                .timeout_global(Some(std::time::Duration::from_secs(10)))
                .build(),
        );

        match agent
            .get(&url)
            .header("Authorization", &format!("Bearer {}", api_key))
            .call()
        {
            Ok(resp) => {
                let body = resp.into_body().read_to_string().unwrap_or_default();
                let parsed: serde_json::Value = serde_json::from_str(&body).unwrap_or_default();
                if let Some(uid) = parsed["zakuro_user_id"].as_str() {
                    let derived = self.config.owner_user_id.as_deref();
                    if derived.is_some() && derived != Some(uid) {
                        // Key-derived ID takes precedence; dashboard mismatch is a warning only.
                        eprintln!("  [HANDSHAKE] owner_user_id mismatch: key-derived={}, dashboard={}. Using key-derived value.",
                                  derived.unwrap_or("?"), uid);
                        return derived.map(str::to_string);
                    }
                    if self.config.owner_user_id.is_none() {
                        self.config.owner_user_id = Some(uid.to_string());
                    }
                    Some(
                        self.config
                            .owner_user_id
                            .clone()
                            .unwrap_or_else(|| uid.to_string()),
                    )
                } else {
                    eprintln!("  [HANDSHAKE] Dashboard did not return zakuro_user_id");
                    None
                }
            }
            Err(e) => {
                eprintln!("  [HANDSHAKE] Failed to verify owner with dashboard: {}", e);
                None
            }
        }
    }
}

pub type SharedBrokerState = Arc<BrokerState>;

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

    #[test]
    fn broker_state_has_shared_async_client() {
        let s = BrokerState::new();
        // Cloning a reqwest::Client is a cheap Arc bump over a shared pool; this
        // asserts the field exists and is usable.
        let _c: reqwest::Client = s.async_http.clone();
    }
}