foxhunt/services/api_gateway/src/main.rs

//! API Gateway Service
//!
//! High-performance gRPC gateway with 6-layer authentication and request routing.
//! Optimized for HFT requirements with <10μs authentication overhead.

use anyhow::Result;
use clap::Parser;
use std::sync::Arc;
use std::time::Duration;
use tracing::{error, info, warn};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

// Import all needed types from the library
use api_gateway::auth::{
        AuditLogger, AuthInterceptor, AuthzService, JwtService, RateLimiter, RevocationService,
    };
use api_gateway::auth::jwt::JwtConfig;

#[derive(Parser, Debug)]
#[command(name = "api_gateway", about = "Foxhunt API Gateway Service")]
struct Args {
    /// gRPC server bind address
    #[arg(long, env = "GATEWAY_BIND_ADDR", default_value = "0.0.0.0:50051")]
    bind_addr: String,

    /// JWT secret (or use JWT_SECRET_FILE for production)
    #[arg(long, env = "JWT_SECRET")]
    jwt_secret: Option<String>,

    /// JWT issuer
    #[arg(long, env = "JWT_ISSUER", default_value = "foxhunt-trading")]
    jwt_issuer: String,

    /// JWT audience
    #[arg(long, env = "JWT_AUDIENCE", default_value = "trading-api")]
    jwt_audience: String,

    /// Redis URL for JWT revocation
    #[arg(long, env = "REDIS_URL", default_value = "redis://localhost:6379")]
    redis_url: String,

    /// Rate limit (requests per second per user)
    #[arg(long, env = "RATE_LIMIT_RPS", default_value = "100")]
    rate_limit_rps: u32,

    /// Enable audit logging
    #[arg(long, env = "ENABLE_AUDIT_LOGGING", default_value = "true")]
    enable_audit_logging: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "api_gateway=info,tower_http=debug".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let args = Args::parse();

    info!("Starting Foxhunt API Gateway Service");
    info!("Bind address: {}", args.bind_addr);
    info!("JWT issuer: {}", args.jwt_issuer);
    info!("JWT audience: {}", args.jwt_audience);
    info!("Redis URL: {}", args.redis_url);
    info!("Rate limit: {} req/s per user", args.rate_limit_rps);
    info!("Audit logging: {}", args.enable_audit_logging);

    // Load JWT configuration (Vault-based with fallback)
    info!("Loading JWT configuration...");
    let jwt_config = match load_jwt_config().await {
        Ok(config) => {
            info!("✅ JWT configuration loaded successfully");
            config
        },
        Err(e) => {
            error!("❌ Failed to load JWT configuration: {}", e);
            return Err(e);
        }
    };

    // Initialize authentication components
    info!("Initializing authentication services...");

    let jwt_service = JwtService::new(jwt_config.jwt_secret.clone(), jwt_config.jwt_issuer.clone(), jwt_config.jwt_audience.clone());
    let jwt_service_rest = JwtService::new(jwt_config.jwt_secret.clone(), jwt_config.jwt_issuer.clone(), jwt_config.jwt_audience.clone());
    info!("✓ JWT service initialized with cached decoding key");

    let revocation_service = RevocationService::new(&args.redis_url)
        .await
        .expect("Failed to connect to Redis for revocation service");
    let revocation_service_rest = RevocationService::new(&args.redis_url)
        .await
        .expect("Failed to connect to Redis for REST API");
    info!("✓ JWT revocation service connected to Redis");

    let authz_service = AuthzService::new();
    info!("✓ Authorization service initialized with permission cache");

    let rate_limiter = RateLimiter::new(args.rate_limit_rps)
        .map_err(|e| anyhow::anyhow!("Failed to create rate limiter: {}", e))?;
    let rate_limiter_rest = RateLimiter::new(args.rate_limit_rps)
        .map_err(|e| anyhow::anyhow!("Failed to create REST rate limiter: {}", e))?;
    info!("✓ Rate limiter initialized ({} req/s)", args.rate_limit_rps);

    let audit_logger = AuditLogger::new(args.enable_audit_logging);
    info!("✓ Audit logger initialized");

    // Create authentication interceptor
    let auth_interceptor = AuthInterceptor::new(
        jwt_service,
        revocation_service,
        authz_service,
        rate_limiter,
        audit_logger,
    );
    info!("✓ 6-layer authentication interceptor ready");

    info!("API Gateway service initialization complete");
    info!("Ready to accept gRPC requests with <10μs auth overhead");

    // Initialize backend service proxies
    info!("Connecting to backend services...");

    let trading_backend_url = std::env::var("TRADING_SERVICE_URL")
        .unwrap_or_else(|_| "http://localhost:50052".to_string());
    let backtesting_backend_url = std::env::var("BACKTESTING_SERVICE_URL")
        .unwrap_or_else(|_| "http://localhost:50053".to_string());
    let ml_training_backend_url = std::env::var("ML_TRAINING_SERVICE_URL")
        .unwrap_or_else(|_| "http://localhost:50054".to_string());

    // Load TLS certificate paths for backtesting service (mTLS)
    let backtesting_tls_ca_cert = std::env::var("BACKTESTING_TLS_CA_CERT").ok();
    let backtesting_tls_client_cert = std::env::var("BACKTESTING_TLS_CLIENT_CERT").ok();
    let backtesting_tls_client_key = std::env::var("BACKTESTING_TLS_CLIENT_KEY").ok();

    // Load TLS certificate paths for ML training service (mTLS)
    let ml_training_tls_ca_cert = std::env::var("ML_TRAINING_TLS_CA_CERT").ok();
    let ml_training_tls_client_cert = std::env::var("ML_TRAINING_TLS_CLIENT_CERT").ok();
    let ml_training_tls_client_key = std::env::var("ML_TRAINING_TLS_CLIENT_KEY").ok();

    // Initialize trading service proxy
    let trading_proxy = api_gateway::grpc::TradingServiceProxy::new_lazy(&trading_backend_url)
        .expect("Failed to create trading service proxy");
    info!("✓ Trading service proxy initialized ({})", trading_backend_url);

    // Initialize backtesting service proxy (optional - graceful degradation)
    info!("Attempting to initialize Backtesting Service proxy...");
    info!("  Backend URL: {}", backtesting_backend_url);
    info!("  CA cert: {:?}", backtesting_tls_ca_cert);
    info!("  Client cert: {:?}", backtesting_tls_client_cert);
    info!("  Client key: {:?}", backtesting_tls_client_key);

    let backtesting_proxy = match api_gateway::grpc::BacktestingServiceProxy::new(
        &backtesting_backend_url,
        backtesting_tls_ca_cert.as_deref(),
        backtesting_tls_client_cert.as_deref(),
        backtesting_tls_client_key.as_deref(),
    ).await {
        Ok(proxy) => {
            info!("✓ Backtesting service proxy initialized ({})", backtesting_backend_url);
            Some(Arc::new(proxy))
        }
        Err(e) => {
            error!("⚠ Backtesting service initialization failed!");
            error!("  Error type: {:?}", e);
            error!("  Error message: {}", e);
            warn!("⚠ Backtesting service unavailable: {}. API Gateway will run without backtesting endpoints.", e);
            warn!("   Backtesting endpoints will return 503 Service Unavailable");
            None
        }
    };

    // Spawn background health check task for backtesting service
    if let Some(proxy) = backtesting_proxy.as_ref() {
        let proxy_clone = Arc::clone(proxy);
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(10));
            loop {
                interval.tick().await;
                proxy_clone.background_health_check().await;
            }
        });
        info!("✓ Backtesting service health check task started (10s interval)");
    }

    // Initialize ML training service proxy (optional - graceful degradation)
    info!("Attempting to initialize ML Training Service proxy...");
    info!("  Backend URL: {}", ml_training_backend_url);
    info!("  CA cert: {:?}", ml_training_tls_ca_cert);
    info!("  Client cert: {:?}", ml_training_tls_client_cert);
    info!("  Client key: {:?}", ml_training_tls_client_key);

    let ml_config = api_gateway::grpc::MlTrainingBackendConfig {
        address: ml_training_backend_url.clone(),
        connect_timeout_ms: 5000,
        request_timeout_ms: 30000,
        circuit_breaker_failures: 5,
        circuit_breaker_reset_secs: 30,
        tls_ca_cert_path: ml_training_tls_ca_cert.clone(),
        tls_client_cert_path: ml_training_tls_client_cert.clone(),
        tls_client_key_path: ml_training_tls_client_key.clone(),
    };
    let ml_training_proxy = match api_gateway::grpc::setup_ml_training_proxy(ml_config).await {
        Ok(proxy) => {
            info!("✓ ML training service proxy initialized ({})", ml_training_backend_url);
            Some(proxy)
        }
        Err(e) => {
            error!("⚠ ML Training service initialization failed!");
            error!("  Error type: {:?}", e);
            error!("  Error message: {}", e);
            warn!("⚠ ML Training service unavailable: {}. API Gateway will run without ML endpoints.", e);
            warn!("   ML training endpoints will return 503 Service Unavailable");
            None
        }
    };

    // Initialize configuration manager (requires database)
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string());
    let db_pool = sqlx::PgPool::connect(&database_url)
        .await
        .expect("Failed to connect to database");
    info!("✓ Database connection established");

    // Create Redis connection for config manager
    let redis_client = redis::Client::open(args.redis_url.clone())
        .expect("Failed to create Redis client");
    let redis_conn = redis::aio::ConnectionManager::new(redis_client)
        .await
        .expect("Failed to create Redis connection manager");

    let mut config_manager = api_gateway::ConfigurationManager::new(db_pool.clone(), redis_conn)
        .await
        .expect("Failed to create configuration manager");

    // Start NOTIFY listener for hot-reload
    config_manager.start_listening()
        .await
        .expect("Failed to start configuration listener");
    info!("✓ Configuration manager initialized with hot-reload");

    // Build gRPC server with all services
    use api_gateway::foxhunt::tli::{
        trading_service_server::TradingServiceServer,
        backtesting_service_server::BacktestingServiceServer,
    };
    use api_gateway::ml_training::ml_training_service_server::MlTrainingServiceServer;

    let addr = args.bind_addr.parse()
        .expect("Failed to parse bind address");

    info!("Starting gRPC server on {}", addr);

    // Setup graceful shutdown
    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();

    // Spawn signal handler for graceful shutdown
    tokio::spawn(async move {
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to listen for ctrl-c");
        info!("Received shutdown signal, draining connections...");
        let _ = shutdown_tx.send(());
    });

    // Setup health check service
    let (health_reporter, health_service) = tonic_health::server::health_reporter();

    // Always register trading service (required)
    health_reporter
        .set_serving::<TradingServiceServer<api_gateway::grpc::TradingServiceProxy>>()
        .await;

    // Conditionally register optional services
    if backtesting_proxy.is_some() {
        health_reporter
            .set_serving::<BacktestingServiceServer<api_gateway::grpc::BacktestingServiceProxy>>()
            .await;
    }
    if ml_training_proxy.is_some() {
        health_reporter
            .set_serving::<MlTrainingServiceServer<api_gateway::grpc::MlTrainingProxy>>()
            .await;
    }

    // Initialize and start Prometheus metrics HTTP endpoint on port 9091
    let gateway_metrics = api_gateway::metrics::GatewayMetrics::new()
        .expect("Failed to initialize gateway metrics");

    // Add service info metric (always present)
    use prometheus::{register_gauge_with_registry, Opts};
    let service_info = register_gauge_with_registry!(
        Opts::new("api_gateway_service_info", "API Gateway service information")
            .const_label("version", env!("CARGO_PKG_VERSION"))
            .const_label("service", "api_gateway"),
        gateway_metrics.registry().as_ref()
    ).expect("Failed to register service info");
    service_info.set(1.0);

    let metrics_registry = gateway_metrics.registry();

    tokio::spawn(async move {
        // Use combined router for metrics AND health/resilience endpoints
        let combined_app = api_gateway::metrics::combined_router(metrics_registry);

        let metrics_addr = "0.0.0.0:9091";
        info!("Prometheus metrics endpoint listening on http://{}", metrics_addr);
        info!("Health endpoints available:");
        info!("  - GET http://{}/health/liveness", metrics_addr);
        info!("  - GET http://{}/health/readiness", metrics_addr);
        info!("  - GET http://{}/health/startup", metrics_addr);
        info!("  - GET http://{}/resilience/circuit-breaker/status", metrics_addr);
        info!("  - GET http://{}/resilience/rate-limit/status", metrics_addr);
        info!("  - GET http://{}/resilience/timeout/config", metrics_addr);
        info!("  - GET http://{}/resilience/retry/config", metrics_addr);

        let listener = tokio::net::TcpListener::bind(metrics_addr)
            .await
            .expect("Failed to bind metrics endpoint");

        axum::serve(listener, combined_app)
            .await
            .expect("Metrics server failed");
    });

    // Initialize REST API server for ML inference endpoints (port 8080)
    if let Some(_ml_proxy) = ml_training_proxy.as_ref() {
        let ml_config_rest = api_gateway::grpc::MlTrainingBackendConfig {
            address: ml_training_backend_url.clone(),
            connect_timeout_ms: 5000,
            request_timeout_ms: 30000,
            circuit_breaker_failures: 5,
            circuit_breaker_reset_secs: 30,
            tls_ca_cert_path: ml_training_tls_ca_cert.clone(),
            tls_client_cert_path: ml_training_tls_client_cert.clone(),
            tls_client_key_path: ml_training_tls_client_key.clone(),
        };
        let ml_client = api_gateway::setup_ml_training_client(ml_config_rest)
            .await
            .expect("Failed to setup ML training client for REST API");

        // Create ML handler state with auth components
        let ml_handler_state = Arc::new(api_gateway::MlHandlerState {
            ml_client,
            auth: Arc::new(auth_interceptor.clone()),
            rate_limiter: Arc::new(rate_limiter_rest),
        });

        // Create auth middleware state (for REST API)
        let auth_middleware_state = Arc::new(api_gateway::AuthMiddlewareState {
            jwt_service: Arc::new(jwt_service_rest),
            revocation_service: Arc::new(revocation_service_rest),
            rate_limiter: ml_handler_state.rate_limiter.clone(),
        });

        // Build ML REST API router with authentication middleware
        use axum::middleware;
        let ml_api_router = api_gateway::ml_router(ml_handler_state)
            .layer(middleware::from_fn_with_state(
                auth_middleware_state,
                api_gateway::jwt_auth_middleware,
            ));

        // Spawn REST API server on port 8080
        tokio::spawn(async move {
            let rest_addr = "0.0.0.0:8080";
            info!("REST API server listening on http://{}", rest_addr);
            info!("ML inference endpoints available:");
            info!("  - POST http://{}/api/v1/ml/predict", rest_addr);
            info!("  - POST http://{}/api/v1/ml/batch_predict", rest_addr);
            info!("  - GET  http://{}/api/v1/ml/model_status", rest_addr);
            info!("  - POST http://{}/api/v1/ml/hot_swap", rest_addr);
            info!("Authentication: JWT Bearer token required (100 req/sec rate limit)");

            let listener = tokio::net::TcpListener::bind(rest_addr)
                .await
                .expect("Failed to bind REST API endpoint");

            axum::serve(listener, ml_api_router)
                .await
                .expect("REST API server failed");
        });
    } else {
        warn!("ML Training Service unavailable - REST API endpoints disabled");
    }

    // Build server with HTTP/2 optimizations
    let mut server_builder = tonic::transport::Server::builder()
        .max_concurrent_streams(Some(10_000))
        .http2_keepalive_interval(Some(Duration::from_secs(30)))
        .http2_keepalive_timeout(Some(Duration::from_secs(10)))
        .layer(tower::ServiceBuilder::new()
            .layer(tower::layer::util::Identity::new())); // Placeholder for auth interceptor layer

    // Add health service
    let mut router = server_builder.add_service(health_service);

    // Always add trading service (required) with authentication
    router = router.add_service(
        TradingServiceServer::with_interceptor(trading_proxy, auth_interceptor.clone())
    );

    // Track service availability for logging
    let backtesting_available = backtesting_proxy.is_some();
    let ml_training_available = ml_training_proxy.is_some();

    // Conditionally add optional services with authentication
    if let Some(backtesting) = backtesting_proxy.as_ref() {
        // Clone the entire Arc - tonic services can work with Arc-wrapped implementations
        router = router.add_service(
            BacktestingServiceServer::with_interceptor(Arc::clone(backtesting), auth_interceptor.clone())
        );
    }
    if let Some(ml_training) = ml_training_proxy {
        router = router.add_service(
            MlTrainingServiceServer::with_interceptor(ml_training, auth_interceptor.clone())
        );
    }

    // Log startup information
    info!("🚀 API Gateway listening on {}", addr);
    info!("   - Trading Service: {} (REQUIRED)", trading_backend_url);
    info!("   - Backtesting Service: {} ({})",
        backtesting_backend_url,
        if backtesting_available { "✓ AVAILABLE" } else { "✗ UNAVAILABLE" });
    info!("   - ML Training Service: {} ({})",
        ml_training_backend_url,
        if ml_training_available { "✓ AVAILABLE" } else { "✗ UNAVAILABLE" });
    info!("   - Health checks: enabled");
    info!("   - Authentication: 6-layer (<10μs overhead)");
    info!("   - Rate limiting: {} req/s per user", args.rate_limit_rps);

    // Start server with graceful shutdown
    let server = router.serve_with_shutdown(addr, async {
        shutdown_rx.await.ok();
    });

    // Start server
    server.await?;

    info!("API Gateway shutdown complete");
    Ok(())
}

/// Load JWT configuration from Vault (production) or environment (development)
///
/// Priority:
/// 1. Vault (secret/foxhunt/jwt) - Production
/// 2. JWT_SECRET_FILE - File-based secret
/// 3. JWT_SECRET env var - Development fallback
async fn load_jwt_config() -> Result<JwtConfig> {
    JwtConfig::new().await
}