iocaine 3.4.0

The deadliest poison known to AI
// SPDX-FileCopyrightText: Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

use anyhow::Result;
use axum::{
    extract::{Query, Request, State},
    handler::Handler,
    response::IntoResponse,
};
use serde::Serialize;
use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;
use tokio_listener::Listener;

use crate::{
    morgue::{Decay, StateOfDecay, shutdown_signal},
    tenx_programmer::TenXProgrammer,
};

use iocaine_powder::{
    acab,
    http::StatusCode,
    sex_dungeon::{DungeonMaster, Language, Request as SDRequest, SharedRequest},
};

#[derive(Clone)]
pub struct Bruce {
    pub state: StateOfDecay,
}

impl Bruce {
    pub fn new(
        language: Language,
        compiler: Option<&impl AsRef<Path>>,
        path: Option<&impl AsRef<Path>>,
        initial_seed: &str,
        metrics: &TenXProgrammer,
        state: &acab::State,
        config: Option<impl Serialize>,
    ) -> Result<Self> {
        let request_handler = DungeonMaster::new(initial_seed)
            .language(language)
            .compiler(compiler.as_ref())
            .path(path.as_ref())
            .config(config)
            .build(&metrics.metrics, state)
            .map_err(|e| anyhow::anyhow!("{e:#?}"))?;
        let request_handler = Arc::new(request_handler);

        if !request_handler.can_output() {
            let path = path
                .as_ref()
                .map_or_else(|| "(default)".into(), |p| p.as_ref().display().to_string());
            tracing::error!({ path }, "At least an output() function is required");
            anyhow::bail!("Requires at least output()");
        }

        let state = Decay {
            metrics: metrics.clone(),
            request_handler,
        }
        .into();

        Ok(Self { state })
    }

    pub async fn serve(self, listener: Listener) -> Result<()> {
        let app = handler
            .layer(tower_http::trace::TraceLayer::new_for_http())
            .with_state(self.state);

        Ok(axum::serve(listener, app.into_make_service())
            .with_graceful_shutdown(shutdown_signal(None))
            .await?)
    }
}

async fn handler(
    State(state): State<StateOfDecay>,
    Query(params): Query<BTreeMap<String, String>>,
    request: Request,
) -> impl IntoResponse {
    let request: SharedRequest = SDRequest {
        method: request.method().to_string(),
        path: request.uri().path().to_owned(),
        headers: request.headers().clone(),
        params,
    }
    .into();

    let state = match state.read() {
        Ok(v) => v,
        Err(e) => {
            tracing::error!("Unable to lock state for reading: {e}");
            return server_error().into_response();
        }
    };
    let decision = if state.request_handler.can_decide() {
        match state.request_handler.decide(request.clone()) {
            Ok(v) => Some(v),
            Err(e) => {
                tracing::error!("Error running decide(): {e}");
                return server_error().into_response();
            }
        }
    } else {
        None
    };
    state.request_handler.output(request, decision).map_or_else(
        |e| {
            tracing::error!("Error running output(): {e}");
            server_error().into_response()
        },
        IntoResponse::into_response,
    )
}

fn server_error() -> impl IntoResponse {
    (StatusCode::INTERNAL_SERVER_ERROR, "")
}