sunbeam-g2v 0.3.3

//! End-to-end example for `sunbeam-g2v` — auth, health, and Connect-RPC features.
//!
//! Demonstrates a realistic service with public routes, JWT-authenticated routes,
//! Keto-authorised routes, a health router wired via `ServerBuilder::with_health`,
//! and a Connect-RPC Eliza service for end-to-end demo with the FE example.
//!
//! # Try it
//!
//! ```text
//! # Start the server (token is printed at startup):
//! cargo run -p sunbeam-g2v --example simple
//!
//! # Public routes (no auth needed):
//! curl http://localhost:8080/
//! curl http://localhost:8080/health/live
//! curl http://localhost:8080/health/ready
//!
//! # Connect-RPC Eliza (no auth needed):
//! curl -X POST -H "Content-Type: application/json" \
//!   -d '{"sentence":"hello"}' \
//!   http://localhost:8080/connectrpc.eliza.v1.ElizaService/Say
//!
//! # Authenticated route — copy the token printed at startup:
//! TOKEN=<paste token here>
//! curl -H "Authorization: Bearer $TOKEN" http://localhost:8080/whoami
//!
//! # Authorised route — requires Keto to grant the tuple first:
//! # (without Keto running, the Keto check returns 500; that is expected)
//! curl -H "Authorization: Bearer $TOKEN" http://localhost:8080/g2v/secrets/42
//!
//! # FE demo:
//! cd libs/sunbeam-g2v-fe/packages/sunbeam-g2v/examples/simple && npm run dev
//! # open http://localhost:5173 — RPC calls will hit this server
//!
//! # Override env:
//! JWT_SECRET=my-secret BIND_ADDR=0.0.0.0:9090 cargo run -p sunbeam-g2v --example simple
//! ```

#![allow(refining_impl_trait_internal, refining_impl_trait_reachable)]

use std::sync::Arc;

use axum::{
    Router as AxumRouter,
    extract::{Path, Request},
    http::StatusCode,
    response::{IntoResponse, Json},
    routing::get,
};
use futures::stream;
use serde_json::json;
use tower_http::cors::{Any, CorsLayer};

use sunbeam_g2v::error::ServiceResult;
use sunbeam_g2v::health::{HealthRouter, KetoHealthCheck};
use sunbeam_g2v::middleware::auth::{
    AuthContext,
    jwt::{JwtLayer, JwtValidator},
    keto::{KetoClient, KetoConfig, KetoLayer},
};
use sunbeam_g2v::router::ServiceRouter;
use sunbeam_g2v::server::{ServerConfig, builder::ServerBuilder};
use sunbeam_g2v::{RequestContext, Router as ConnectRouter};

// Pull in the generated Eliza service trait, message types, and the Ext trait.
// Wrapped in a module to avoid name collision with the `connectrpc` extern crate.
mod eliza_proto {
    include!(concat!(env!("OUT_DIR"), "/_eliza.rs"));
}

use eliza_proto::connectrpc::eliza::v1::{
    ElizaService, ElizaServiceExt, IntroduceRequest, IntroduceResponse, SayRequest, SayResponse,
};

// ============================================================================
// Eliza service implementation
// ============================================================================

struct ElizaServiceImpl;

impl ElizaService for ElizaServiceImpl {
    async fn say(
        &self,
        _ctx: RequestContext,
        request: connectrpc::ServiceRequest<'_, SayRequest>,
    ) -> connectrpc::ServiceResult<SayResponse> {
        let sentence = request.view().sentence.to_lowercase();
        let reply = if sentence.contains("hello") || sentence.contains("hi") {
            "Hello! I'm Eliza, your digital therapist. How are you feeling today?"
        } else if sentence.contains("feel") || sentence.contains("feeling") {
            "Tell me more about how you're feeling."
        } else if sentence.contains("sad")
            || sentence.contains("unhappy")
            || sentence.contains("depressed")
        {
            "I'm sorry to hear that. What do you think is causing these feelings?"
        } else if sentence.contains("happy")
            || sentence.contains("good")
            || sentence.contains("great")
        {
            "I'm glad to hear that! What's been making you feel this way?"
        } else if sentence.contains("bye") || sentence.contains("goodbye") {
            "Goodbye! Take care of yourself."
        } else if sentence.contains("?") {
            "That's an interesting question. What do you think the answer is?"
        } else if sentence.is_empty() {
            "I'm listening. Please, go on."
        } else {
            "Please, tell me more."
        };

        Ok(connectrpc::Response::new(SayResponse {
            sentence: reply.to_string(),
            ..Default::default()
        }))
    }

    async fn introduce(
        &self,
        _ctx: RequestContext,
        request: connectrpc::ServiceRequest<'_, IntroduceRequest>,
    ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<IntroduceResponse>> {
        let name = request.view().name.to_string();
        let responses = vec![
            Ok(IntroduceResponse {
                sentence: format!("Hi {name}! I'm Eliza, a digital therapist."),
                ..Default::default()
            }),
            Ok(IntroduceResponse {
                sentence: "I'm here to listen and help you explore your thoughts.".to_string(),
                ..Default::default()
            }),
            Ok(IntroduceResponse {
                sentence: "Feel free to tell me what's on your mind.".to_string(),
                ..Default::default()
            }),
        ];
        connectrpc::Response::stream_ok(stream::iter(responses))
    }
}

// ============================================================================
// Axum route handlers
// ============================================================================

/// `GET /` — public, no auth required.
async fn index_handler() -> impl IntoResponse {
    Json(json!({ "service": "sunbeam-g2v simple example", "status": "ok" }))
}

/// `GET /whoami` — requires a valid JWT.
async fn whoami_handler(request: Request) -> impl IntoResponse {
    let ctx = request.extensions().get::<AuthContext>().cloned();
    match ctx {
        Some(ctx) if ctx.is_authenticated() => {
            let sub = ctx
                .subject()
                .cloned()
                .unwrap_or_else(|| "<unknown>".to_string());
            Json(json!({ "subject": sub })).into_response()
        }
        _ => (
            StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "unauthenticated" })),
        )
            .into_response(),
    }
}

/// `GET /g2v/secrets/:id` — requires JWT then Keto `read` on `G2vTest`.
async fn secret_handler(Path(id): Path<String>, request: Request) -> impl IntoResponse {
    let ctx = request.extensions().get::<AuthContext>().cloned();
    let sub = ctx
        .and_then(|c| c.subject().cloned())
        .unwrap_or_else(|| "<unknown>".to_string());
    Json(json!({
        "subject": sub,
        "secret_id": id,
        "message": "access granted"
    }))
}

// ============================================================================
// Main
// ============================================================================

#[tokio::main]
async fn main() -> ServiceResult<()> {
    // ---- configuration from env (no panics on missing vars) ----------------
    let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "dev-secret".to_string());

    let keto_grpc_url =
        std::env::var("KETO_GRPC_URL").unwrap_or_else(|_| "http://localhost:4466".to_string());

    let keto_write_grpc_url = std::env::var("KETO_WRITE_GRPC_URL")
        .unwrap_or_else(|_| "http://localhost:4467".to_string());

    let bind_addr: std::net::SocketAddr = std::env::var("BIND_ADDR")
        .unwrap_or_else(|_| "127.0.0.1:8080".to_string())
        .parse()
        .unwrap_or_else(|_| "127.0.0.1:8080".parse().unwrap());

    // ---- shared state -------------------------------------------------------
    let validator = Arc::new(JwtValidator::with_secret(&jwt_secret));
    let keto_client = Arc::new(KetoClient::new(KetoConfig {
        grpc_endpoint: keto_grpc_url,
        write_grpc_endpoint: keto_write_grpc_url,
    }));

    // ---- CORS (permissive for dev — allows the FE at localhost:5173) --------
    // Applied as a layer on all axum routes so preflight OPTIONS requests succeed
    // for both RPC paths (/connectrpc.eliza.v1.ElizaService/Say) and REST routes.
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any)
        .expose_headers(Any);

    // ---- Connect-RPC Eliza service -----------------------------------------
    // Register the Eliza service on the ConnectRouter, then wrap it in a
    // ServiceRouter so the ServerBuilder can own it cleanly (it becomes the
    // single fallback via into_axum_router() inside AxumServer).
    let eliza = Arc::new(ElizaServiceImpl);
    let connect_router: ConnectRouter = eliza.register(ConnectRouter::new());
    let service_router = ServiceRouter::from_router(connect_router);

    // ---- Axum Routes -------------------------------------------------------

    let public_routes = AxumRouter::new().route("/", get(index_handler));

    let authed_routes = AxumRouter::new()
        .route("/whoami", get(whoami_handler))
        .layer(JwtLayer::new((*validator).clone()));

    let keto_layer = KetoLayer::new((*keto_client).clone(), "G2vTest", "read")
        .with_object_extractor(std::sync::Arc::new(|req: &http::Request<()>| {
            req.uri()
                .path()
                .rsplit('/')
                .next()
                .unwrap_or("")
                .to_string()
        }));

    let authzd_routes = AxumRouter::new()
        .route("/g2v/secrets/{id}", get(secret_handler))
        .layer(keto_layer)
        .layer(JwtLayer::new((*validator).clone()));

    let app_routes = public_routes.merge(authed_routes).merge(authzd_routes);

    // ---- Health router -----------------------------------------------------
    let health =
        HealthRouter::new().with_check(Arc::new(KetoHealthCheck::new((*keto_client).clone())));

    // ---- Server config ------------------------------------------------------
    let config = ServerConfig {
        addr: bind_addr,
        name: "simple-example".to_string(),
        ..ServerConfig::default()
    };

    // ---- Dev-mode token hint ------------------------------------------------
    let sample_token = validator
        .create_user_token("dev-user")
        .unwrap_or_else(|_| "<token-mint-failed>".to_string());

    println!("listening on http://{}", config.addr);
    println!();
    println!("  sample token (dev-user):");
    println!("    {sample_token}");
    println!();
    println!("  curl http://{}/", config.addr);
    println!(
        "  curl -H 'Authorization: Bearer {sample_token}' http://{}/whoami",
        config.addr
    );
    println!(
        "  curl -H 'Authorization: Bearer {sample_token}' http://{}/g2v/secrets/42",
        config.addr
    );
    println!("  curl http://{}/health/live", config.addr);
    println!("  curl http://{}/health/ready", config.addr);
    println!();
    println!("  # Connect-RPC Eliza (no auth):");
    println!(
        "  curl -X POST -H 'Content-Type: application/json' -d '{{\"sentence\":\"hello\"}}' \\",
    );
    println!(
        "    http://{}/connectrpc.eliza.v1.ElizaService/Say",
        config.addr
    );
    println!();

    // ---- Graceful shutdown --------------------------------------------------
    let shutdown = async {
        let _ = tokio::signal::ctrl_c().await;
        println!("shutdown signal received");
    };

    // Build the full axum app (connect service as fallback + extra routes merged in),
    // then apply CORS as the outermost layer so preflight OPTIONS on the RPC path
    // (/connectrpc.eliza.v1.ElizaService/Say) is handled before the Connect handler
    // rejects it with 415.
    let server = ServerBuilder::new()
        .with_router(service_router)
        .with_config(config)
        .with_health(health)
        .with_routes(app_routes)
        .build_axum()?;

    let app = server.app().layer(cors);

    let listener = tokio::net::TcpListener::bind(server.config().addr)
        .await
        .map_err(|e| sunbeam_g2v::error::ServiceError::Internal(format!("bind: {e}")))?;

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown)
        .await
        .map_err(|e| sunbeam_g2v::error::ServiceError::Internal(format!("axum::serve: {e}")))
}