zc2 0.0.14

//! Dashboard API-based central ledger for credit management.
//!
//! This module provides a secure, centralized credit system:
//! - Credits managed via dashboard API (single source of truth)
//! - API key authentication via key format extraction
//! - Atomic transactions via dashboard API
//! - Full transaction history via API

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// Transaction record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LedgerTransaction {
    pub id: String,
    pub user_id: String,
    pub tx_type: TransactionType,
    pub amount: f64,
    pub balance_after: f64,
    pub timestamp: DateTime<Utc>,
    pub reference: Option<String>,
    pub authorized_by: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum TransactionType {
    Credit,
    Debit,
    Reserve,
    Commit,
    Cancel,
    Refund,
}

/// Ledger error types
#[derive(Debug)]
pub enum LedgerError {
    ConnectionFailed(String),
    Unauthorized(String),
    InsufficientCredits { required: f64, available: f64 },
    InvalidReservation(String),
    DatabaseError(String),
}

impl std::fmt::Display for LedgerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ConnectionFailed(msg) => write!(f, "Ledger connection failed: {}", msg),
            Self::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
            Self::InsufficientCredits { required, available } => {
                write!(f, "Insufficient credits: need {}, have {}", required, available)
            }
            Self::InvalidReservation(id) => write!(f, "Invalid reservation: {}", id),
            Self::DatabaseError(msg) => write!(f, "Database error: {}", msg),
        }
    }
}

/// Central ledger for credit management
pub struct Ledger {
    api_url: Option<String>,
    api_key: Option<String>,
    /// reservation_id → (user_id, amount)
    local_reservations: dashmap::DashMap<String, (String, f64)>,
    /// Local in-memory credits (used when API unavailable or for local ops)
    pub local_credits: dashmap::DashMap<String, f64>,
    /// Authoritative in-memory balances for P2P mode.
    /// Only populated for users this broker is authoritative for.
    pub authoritative_balances: dashmap::DashMap<String, f64>,
}

impl Ledger {
    /// Create a new ledger with the given api_url and api_key.
    pub fn new(api_url: Option<String>, api_key: Option<String>) -> Self {
        Self {
            api_url,
            api_key,
            local_reservations: dashmap::DashMap::new(),
            local_credits: dashmap::DashMap::new(),
            authoritative_balances: dashmap::DashMap::new(),
        }
    }

    /// Returns true when the broker is in API mode (dashboard API is the billing source).
    ///
    /// Standalone mode takes priority: if `ZAKURO_MASTER_KEY` is set, local credits
    /// are the billing source regardless of whether `api_url` is also configured.
    /// `api_url` may still be present for tx_buffer publishing even in standalone mode.
    pub fn is_api_mode(&self) -> bool {
        let master_key = std::env::var("ZAKURO_MASTER_KEY").unwrap_or_default();
        if !master_key.is_empty() {
            return false; // standalone billing always wins
        }
        self.api_url.is_some() && self.api_key.is_some()
    }

    /// Resolve user_id from an API key (Bearer token).
    /// Returns the zakuro_user_id associated with the key, or an error.
    pub fn resolve_user_from_api_key(&self, api_key: &str) -> Result<String, LedgerError> {
        Self::extract_user_from_key_format(api_key)
    }

    /// Extract user_id from key format: zk_{user_id}_{random_hex}
    /// ZAKURO_MASTER_KEY env var resolves to "admin". Any non-zk_ key also
    /// resolves to "admin" when ZAKURO_MASTER_KEY is not configured (standalone/dev mode).
    fn extract_user_from_key_format(api_key: &str) -> Result<String, LedgerError> {
        if api_key.is_empty() {
            return Err(LedgerError::Unauthorized("Empty API key".to_string()));
        }
        // Check master key from env → return "admin"
        let master_key = std::env::var("ZAKURO_MASTER_KEY").unwrap_or_default();
        if !master_key.is_empty() && api_key == master_key {
            return Ok("admin".to_string());
        }
        // Standard zk_ format: zk_{user_id}_{random_hex}
        if let Some(rest) = api_key.strip_prefix("zk_") {
            if let Some(pos) = rest.rfind('_') {
                let user_id = &rest[..pos];
                if !user_id.is_empty() {
                    return Ok(user_id.to_string());
                }
            }
            return Err(LedgerError::Unauthorized("Invalid zk_ key format".to_string()));
        }
        // In standalone mode (no ZAKURO_MASTER_KEY set), any non-zk_ key
        // is treated as an admin token (dev/test environments).
        if master_key.is_empty() {
            return Ok("admin".to_string());
        }
        Err(LedgerError::Unauthorized("Invalid API key format".to_string()))
    }

    /// Get user's current balance — tries dashboard API first, falls back to local cache.
    pub fn get_balance(&self, user_id: &str) -> f64 {
        // In standalone mode (MASTER_KEY set), skip dashboard API — local_credits is authoritative.
        let master_key = std::env::var("ZAKURO_MASTER_KEY").unwrap_or_default();
        // Try dashboard API first (only when not in standalone mode)
        if master_key.is_empty() {
            if let (Some(ref api_url), Some(ref api_key)) = (&self.api_url, &self.api_key) {
                let endpoint = format!("{}/api/broker/balance/{}", api_url.trim_end_matches('/'), user_id);
                if let Ok(resp) = ureq::get(&endpoint)
                    .set("X-Broker-Api-Key", api_key)
                    .call()
                {
                    if resp.status() == 200 {
                        let body = resp.into_string().unwrap_or_default();
                        let parsed: serde_json::Value = serde_json::from_str(&body).unwrap_or_default();
                        // Dashboard API returns "credits_balance"; fallback accepts "balance" too
                        let balance_val = parsed["credits_balance"].as_f64()
                            .or_else(|| parsed["balance"].as_f64());
                        if let Some(balance) = balance_val {
                            // Populate BOTH credit maps so reserve() works correctly.
                            // local_credits is consumed by reserve()/commit() (Standalone path).
                            // authoritative_balances is used by local_reserve()/local_commit() (P2P Local path).
                            self.local_credits.insert(user_id.to_string(), balance);
                            self.authoritative_balances.insert(user_id.to_string(), balance);
                            return balance;
                        }
                    }
                }
            }
        }
        // Fall back to whichever local cache has data
        self.authoritative_balances
            .get(user_id).map(|v| *v)
            .or_else(|| self.local_credits.get(user_id).map(|v| *v))
            .unwrap_or(0.0)
    }

    /// Reserve credits for a pending operation (local in-memory).
    ///
    /// Requires that `get_balance()` has been called first to populate local_credits.
    /// If local_credits is not yet populated for this user (race condition or first call),
    /// falls back to authoritative_balances as a safety net.
    pub fn reserve(
        &self,
        user_id: &str,
        amount: f64,
        _reference: &str,
    ) -> Result<String, LedgerError> {
        let reservation_id = uuid::Uuid::new_v4().to_string();

        // If local_credits is not yet seeded for this user (shouldn't happen after get_balance fix,
        // but guard defensively), pull from authoritative_balances.
        if !self.local_credits.contains_key(user_id) {
            if let Some(bal) = self.authoritative_balances.get(user_id).map(|v| *v) {
                self.local_credits.insert(user_id.to_string(), bal);
            }
        }

        let mut success = false;
        self.local_credits
            .entry(user_id.to_string())
            .and_modify(|b| {
                if *b >= amount {
                    *b -= amount;
                    success = true;
                }
            });
        if !success {
            let balance = self.local_credits.get(user_id).map(|v| *v).unwrap_or(0.0);
            return Err(LedgerError::InsufficientCredits { required: amount, available: balance });
        }
        // Keep authoritative_balances in sync so /me and P2P billing see the correct balance.
        let balance_after = self.local_credits.get(user_id).map(|v| *v).unwrap_or(0.0);
        self.authoritative_balances.entry(user_id.to_string())
            .and_modify(|b| *b = balance_after)
            .or_insert(balance_after);
        self.local_reservations.insert(reservation_id.clone(), (user_id.to_string(), amount));
        Ok(reservation_id)
    }

    /// Commit a reservation (finalize the charge, refund difference).
    pub fn commit(&self, reservation_id: &str, actual_amount: f64) -> Result<f64, LedgerError> {
        if let Some((_, (user_id, reserved))) = self.local_reservations.remove(reservation_id) {
            let refund = reserved - actual_amount;
            if refund > 0.0 {
                self.local_credits.entry(user_id.clone()).and_modify(|b| *b += refund);
            }
            let final_bal = self.local_credits.get(&user_id).map(|v| *v).unwrap_or(0.0);
            // Keep authoritative_balances in sync so /me reflects the final balance.
            self.authoritative_balances.entry(user_id.to_string())
                .and_modify(|b| *b = final_bal)
                .or_insert(final_bal);
            return Ok(final_bal);
        }
        Err(LedgerError::InvalidReservation(reservation_id.to_string()))
    }

    /// Cancel a reservation (refund the full amount).
    pub fn cancel(&self, reservation_id: &str) -> Result<(), LedgerError> {
        if let Some((_, (user_id, amount))) = self.local_reservations.remove(reservation_id) {
            if amount > 0.0 {
                self.local_credits.entry(user_id).and_modify(|b| *b += amount);
            }
        }
        Ok(())
    }

    /// Cancel a reservation from WAL replay (no DashMap entry — refund directly).
    pub fn cancel_from_wal(&self, user_id: &str, amount: f64) -> Result<(), LedgerError> {
        if amount <= 0.0 {
            return Ok(());
        }
        self.local_credits.entry(user_id.to_string()).and_modify(|b| *b += amount).or_insert(amount);
        self.authoritative_balances.entry(user_id.to_string()).and_modify(|b| *b += amount);
        Ok(())
    }

    /// Commit from WAL replay (no DashMap entry — refund difference directly).
    pub fn commit_from_wal(&self, user_id: &str, reserved: f64, actual: f64) -> Result<f64, LedgerError> {
        let refund = reserved - actual;
        if refund > 0.0 {
            self.local_credits.entry(user_id.to_string()).and_modify(|b| *b += refund).or_insert(refund);
            self.authoritative_balances.entry(user_id.to_string()).and_modify(|b| *b += refund);
        }
        Ok(self.get_balance(user_id))
    }

    /// Publish a transaction event via the dashboard API.
    pub fn publish_transaction(
        &self,
        _request_id: &str,
        user_id: &str,
        tx_type: &str,
        amount: f64,
        _balance_after: f64,
        worker_id: &str,
        duration_ms: f64,
        source_node: Option<&str>,
    ) {
        let (job_name, dashboard_type, status, credits_amount, compute_hours) = match tx_type {
            "commit" => {
                let job = if worker_id.is_empty() {
                    "Compute Job".to_string()
                } else {
                    format!("Compute Job ({})", worker_id)
                };
                let hours = if duration_ms > 0.0 { duration_ms / 3_600_000.0 } else { 0.0 };
                (job, "job_execution", "completed", amount, hours)
            }
            "credit" => {
                ("zkcr Added (Broker)".to_string(), "credit_purchase", "completed", amount, 0.0)
            }
            "cancel" => {
                let job = if worker_id.is_empty() {
                    "Cancelled Job".to_string()
                } else {
                    format!("Cancelled Job ({})", worker_id)
                };
                (job, "job_execution", "failed", 0.0, 0.0)
            }
            _ => return,
        };

        let compute_hours_opt: Option<f64> = if compute_hours > 0.0 { Some(compute_hours) } else { None };
        let duration_opt: Option<f64> = if duration_ms > 0.0 { Some(duration_ms) } else { None };

        if let (Some(ref api_url), Some(ref api_key)) = (&self.api_url, &self.api_key) {
            let worker_id_opt: Option<&str> = if worker_id.is_empty() { None } else { Some(worker_id) };
            let payload = serde_json::json!({
                "zakuro_user_id": user_id,
                "job_name": job_name,
                "transaction_type": dashboard_type,
                "credits_amount": credits_amount,
                "status": status,
                "compute_hours": compute_hours_opt,
                "worker_id": worker_id_opt,
                "source_node": source_node,
                "metadata": null,
                "duration_ms": duration_opt,
            });
            let endpoint = format!("{}/api/broker/transaction", api_url.trim_end_matches('/'));
            let result = ureq::post(&endpoint)
                .set("X-Broker-Api-Key", api_key)
                .set("Content-Type", "application/json")
                .send_string(&serde_json::to_string(&payload).unwrap_or_default());
            if let Err(e) = result {
                eprintln!("  [LEDGER] Failed to POST transaction via API: {}", e);
            }
        }
    }

    /// Get user info including balance.
    pub fn get_user_info(&self, user_id: &str) -> UserInfo {
        let balance = self.get_balance(user_id);
        UserInfo {
            user_id: user_id.to_string(),
            balance,
        }
    }

    /// Sync workers via Dashboard API instead of direct PostgreSQL access.
    /// This is the only supported worker sync method.
    ///
    /// `broker_tailscale_ip` is the broker's own Tailscale IP. It is used as the
    /// `tailscale_ip` for workers whose URI resolves to a loopback address (Docker
    /// workers share the broker's network namespace and have no independent IP).
    pub fn sync_workers_via_api(
        zakuro_user_id: &str,
        workers: &[super::worker::Worker],
        api_url: &str,
        api_key: &str,
        node_name: Option<&str>,
        broker_tailscale_ip: Option<&str>,
    ) -> Result<(), String> {
        use serde_json::json;

        if workers.is_empty() {
            return Ok(());
        }

        // Build worker sync payloads
        let worker_payloads: Vec<_> = workers
            .iter()
            .map(|worker| {
                let status_str = match worker.status {
                    super::worker::WorkerStatus::Healthy => "online",
                    super::worker::WorkerStatus::Busy => "online",
                    super::worker::WorkerStatus::Unhealthy => "offline",
                    super::worker::WorkerStatus::Draining => "offline",
                };

                let cpu_cores = worker.resources.cpus_available as i32;
                let ram_gb = (worker.resources.memory_available as f64 / (1024.0 * 1024.0 * 1024.0)).round() as i32;

                // Docker workers share the broker's network namespace, so their URI IP
                // is loopback. Use the broker's own Tailscale IP instead.
                let effective_tailscale_ip = match worker.tailscale_ip.as_deref() {
                    Some("127.0.0.1") | Some("::1") | Some("localhost") | None => broker_tailscale_ip,
                    Some(ip) => Some(ip),
                };

                let price_per_hour = worker.pricing.price_per_hour;

                json!({
                    "zakuro_user_id": zakuro_user_id,
                    "worker_id": &worker.name,  // Use stable name as worker_id
                    "name": &worker.name,
                    "status": status_str,
                    "gpu_model": worker.hardware.gpu_model.as_deref(),
                    "gpu_vram_gb": worker.hardware.gpu_vram_gb.map(|v| v as i32),
                    "cpu_model": worker.hardware.cpu_model.as_deref(),
                    "cpu_cores": cpu_cores,
                    "ram_gb": ram_gb,
                    "storage_gb": worker.hardware.storage_gb.map(|v| v as i32),
                    "source_node": node_name,
                    "tailscale_ip": effective_tailscale_ip,
                    "is_docker": worker.is_docker,
                    "price_per_hour": price_per_hour,
                    "min_charge": worker.pricing.min_charge,
                })
            })
            .collect();

        // Call batch sync API
        let endpoint = format!(
            "{}/api/broker/sync-workers?zakuro_user_id={}",
            api_url.trim_end_matches('/'),
            zakuro_user_id
        );

        // Body is just the list of workers
        let payload_str = serde_json::to_string(&worker_payloads)
            .map_err(|e| format!("Failed to serialize payload: {}", e))?;

        let response = ureq::post(&endpoint)
            .set("X-Broker-Api-Key", api_key)
            .set("Content-Type", "application/json")
            .send_string(&payload_str);

        match response {
            Ok(resp) => {
                if resp.status() == 200 {
                    Ok(())
                } else {
                    Err(format!("API returned status {}", resp.status()))
                }
            }
            Err(e) => Err(format!("Failed to sync workers via API: {}", e)),
        }
    }
}

/// User info response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserInfo {
    pub user_id: String,
    pub balance: f64,
}

// --- P2P local credit operations (zero API calls on hot path) ---

impl Ledger {
    /// Load a user's balance from the dashboard API if not already cached in authoritative_balances.
    /// Returns the cached balance.
    pub fn load_balance_if_needed(&self, user_id: &str) -> f64 {
        // In standalone mode (no dashboard, or MASTER_KEY set), local_credits is the source of truth.
        // Sync it to authoritative_balances so the P2P billing path sees the correct balance.
        let standalone = self.api_url.is_none()
            || !std::env::var("ZAKURO_MASTER_KEY").unwrap_or_default().is_empty();
        if standalone {
            if let Some(bal) = self.local_credits.get(user_id).map(|v| *v) {
                self.authoritative_balances.insert(user_id.to_string(), bal);
                return bal;
            }
        }

        if let Some(balance) = self.authoritative_balances.get(user_id) {
            return *balance;
        }

        // First access — load from dashboard API
        let balance = self.get_balance(user_id);
        self.authoritative_balances.insert(user_id.to_string(), balance);
        balance
    }

    /// Get the authoritative in-memory balance (returns None if not cached).
    pub fn get_authoritative_balance(&self, user_id: &str) -> Option<f64> {
        self.authoritative_balances.get(user_id).map(|v| *v)
    }

    /// Reserve credits in-memory (P2P local path — zero API calls).
    /// Returns (reservation_id, balance_before) on success.
    pub fn local_reserve(
        &self,
        user_id: &str,
        amount: f64,
        _reference: &str,
    ) -> Result<(String, f64), LedgerError> {
        let reservation_id = uuid::Uuid::new_v4().to_string();

        // Ensure balance is loaded
        self.load_balance_if_needed(user_id);

        // Atomic deduct from DashMap
        let mut success = false;
        let mut balance_before = 0.0;
        self.authoritative_balances
            .entry(user_id.to_string())
            .and_modify(|b| {
                balance_before = *b;
                if *b >= amount {
                    *b -= amount;
                    success = true;
                }
            });

        if !success {
            let available = self.authoritative_balances
                .get(user_id)
                .map(|v| *v)
                .unwrap_or(0.0);
            return Err(LedgerError::InsufficientCredits {
                required: amount,
                available,
            });
        }

        // Track reservation for commit/cancel
        self.local_reservations.insert(
            reservation_id.clone(),
            (user_id.to_string(), amount),
        );

        Ok((reservation_id, balance_before))
    }

    /// Commit a reservation in-memory (P2P local path — zero API calls).
    /// Returns balance_after.
    pub fn local_commit(
        &self,
        reservation_id: &str,
        actual_cost: f64,
    ) -> Result<f64, LedgerError> {
        if let Some((_, (user_id, reserved))) = self.local_reservations.remove(reservation_id) {
            let refund = reserved - actual_cost;
            if refund > 0.0 {
                self.authoritative_balances
                    .entry(user_id.clone())
                    .and_modify(|b| *b += refund);
            }
            let balance_after = self.authoritative_balances
                .get(&user_id)
                .map(|v| *v)
                .unwrap_or(0.0);
            Ok(balance_after)
        } else {
            Err(LedgerError::InvalidReservation(reservation_id.to_string()))
        }
    }

    /// Cancel a reservation in-memory (P2P local path — zero API calls).
    pub fn local_cancel(&self, reservation_id: &str) -> Result<(), LedgerError> {
        if let Some((_, (user_id, amount))) = self.local_reservations.remove(reservation_id) {
            if amount > 0.0 {
                self.authoritative_balances
                    .entry(user_id)
                    .and_modify(|b| *b += amount);
            }
        }
        // Idempotent — OK even if not found
        Ok(())
    }

    /// Add credits to a user's authoritative in-memory balance (P2P earn path).
    /// Loads balance from dashboard API first if not already cached.
    /// Returns the new balance after the addition.
    pub fn local_add_credits(&self, user_id: &str, amount: f64) -> f64 {
        self.load_balance_if_needed(user_id);
        self.authoritative_balances
            .entry(user_id.to_string())
            .and_modify(|b| *b += amount)
            .or_insert(amount);
        self.authoritative_balances
            .get(user_id)
            .map(|v| *v)
            .unwrap_or(amount)
    }

    /// Get all authoritative balance entries (for flush thread).
    pub fn authoritative_balance_snapshot(&self) -> Vec<(String, f64)> {
        self.authoritative_balances
            .iter()
            .map(|e| (e.key().clone(), *e.value()))
            .collect()
    }
}

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

    #[test]
    fn test_local_fallback() {
        let ledger = Ledger::new(None, None);
        // Seed local credits manually
        ledger.local_credits.insert("user1".to_string(), 100.0);

        // Reserve
        let reservation = ledger.reserve("user1", 10.0, "test");
        assert!(reservation.is_ok());
        assert_eq!(ledger.local_credits.get("user1").map(|v| *v).unwrap_or(0.0), 90.0);

        // Commit partial
        let commit = ledger.commit(&reservation.unwrap(), 5.0);
        assert!(commit.is_ok());
        assert_eq!(ledger.local_credits.get("user1").map(|v| *v).unwrap_or(0.0), 95.0);
    }

    #[test]
    fn test_extract_user_from_key() {
        assert_eq!(
            Ledger::extract_user_from_key_format("zk_9000000001_abc123def456").unwrap(),
            "9000000001"
        );
        // In standalone mode (no ZAKURO_MASTER_KEY), any non-zk_ key resolves to "admin"
        // (dev/test environment convenience). An error is only returned when ZAKURO_MASTER_KEY
        // is set and the key doesn't match it.
        if std::env::var("ZAKURO_MASTER_KEY").map(|v| !v.is_empty()).unwrap_or(false) {
            assert!(Ledger::extract_user_from_key_format("invalid").is_err());
        } else {
            assert_eq!(
                Ledger::extract_user_from_key_format("invalid").unwrap(),
                "admin"
            );
        }
    }
}