zc2 0.0.12

//! 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
//! - Tailscale-based P2P worker discovery
//! - Live terminal UI for monitoring
//! - Benchmarking tools

pub mod bench;
pub mod config;
pub mod uri;
pub mod credits;
pub mod discovery;
pub mod flush;
pub mod info;
pub mod ledger;
pub mod peer;
pub mod quic;
pub mod recovery;
pub mod router;
pub mod server;
pub mod stats;
pub mod task_board;
pub mod tui;
pub mod wal;
pub mod worker;
pub mod worker_quic;

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

use dashmap::DashMap;

pub use config::{fetch_broker_config, get_tailscale_auth_key};
pub use credits::{CreditManager, UserCredits};
pub use discovery::{detect_discovery_mode, Discovery, DiscoveryConfig, DiscoveryMode};
pub use ledger::{Ledger, LedgerError};
pub use router::{Router, RoutingStrategy};
pub use server::start_server;
pub use stats::{StatsCollector, StatsResponse, WorkerStats};
pub use tui::{cleanup_terminal, run_remote_tui, run_tui};
pub use flush::TransactionBuffer;
pub use peer::PeerManager;
pub use wal::Wal;
pub use worker::{HardwareInfo, Worker, WorkerRegistry, WorkerStatus};

/// Attempt to infer the node name from Tailscale, falling back to the system hostname.
///
/// Priority:
/// 1. `tailscale status --json` subprocess (works when tailscale CLI is on PATH)
/// 2. Tailscale local API socket at `/var/run/tailscale/tailscaled.sock`
/// 3. System hostname (`hostname` command)
fn detect_node_name() -> Option<String> {
    // 1. Try tailscale CLI
    if let Ok(out) = std::process::Command::new("tailscale")
        .args(["status", "--json"])
        .output()
    {
        if out.status.success() {
            if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&out.stdout) {
                if let Some(name) = v["Self"]["HostName"].as_str().filter(|s| !s.is_empty()) {
                    return Some(name.to_string());
                }
            }
        }
    }

    // 2. Try Tailscale local socket
    // Use curl as a subprocess since we don't pull in a Unix-socket HTTP client here
    if let Ok(out) = std::process::Command::new("curl")
        .args([
            "--silent",
            "--unix-socket", "/var/run/tailscale/tailscaled.sock",
            "http://local-tailscaled.sock/localapi/v0/status",
        ])
        .output()
    {
        if out.status.success() {
            if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&out.stdout) {
                if let Some(name) = v["Self"]["HostName"].as_str().filter(|s| !s.is_empty()) {
                    return Some(name.to_string());
                }
            }
        }
    }

    // 3. Fall back to system hostname
    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())
}

/// 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 TUI mode (interactive terminal dashboard)
    pub tui_mode: bool,
    /// Enable Tailscale-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 Tailscale IP detection (test-only, avoids global env var race).
    /// When set, `BrokerState::own_tailscale_ip` uses this value instead of
    /// probing network interfaces.
    pub tailscale_ip_override: Option<String>,
    /// Explicit QUIC port (UDP).  When `None`, defaults to `port + 1`.
    pub quic_port: Option<u16>,
}

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,
            tui_mode: false,
            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),
            peer_key: std::env::var("ZAKURO_PEER_KEY").ok(),
            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(),
            tailscale_ip_override: None,
            quic_port: None,
        }
    }
}

/// 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 for TUI
    pub stats: Arc<StatsCollector>,
    /// Write-ahead log for crash recovery
    pub wal: Wal,
    /// This node's own Tailscale IP (for detecting local vs remote workers)
    pub own_tailscale_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,
}

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_tailscale_ip = config.tailscale_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 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());
            }
            discovered
        } else {
            config.discovery.peers.clone()
        };

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

        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_tailscale_ip,
            peer_manager,
            tx_buffer: TransactionBuffer::new(),
            quic: OnceLock::new(),
            quic_port: std::sync::atomic::AtomicU16::new(0),
        }
    }

    /// 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_tailscale_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::AgentBuilder::new()
            .timeout(std::time::Duration::from_secs(10))
            .build();

        match agent.get(&url)
            .set("Authorization", &format!("Bearer {}", api_key))
            .call()
        {
            Ok(resp) => {
                let body = resp.into_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>;