//! Required env: SUBTOKENLING_LISTEN, SUBTOKENLING_DB, SUBTOKENLING_MODEL_DIR,
//! SUBTOKENLING_MODEL_ID, SUBTOKENLING_ALLOWED_ORIGIN, ORT_DYLIB_PATH.
//! Optional: SUBTOKENLING_MODEL_DTYPE=q4|q8|fp32 (default q4); SUBTOKENLING_THREADS=4;
//! SUBTOKENLING_GPU_DEVICE (logical CUDA index after CUDA_VISIBLE_DEVICES,
//! requires a CUDA-enabled ORT library); SUBTOKENLING_MAX_PETS=64;
//! SUBTOKENLING_WEATHER_SEED=""; SUBTOKENLING_FEATURES="[]" (JSON feature names).
//! No state, access keys, prompts or generated text are logged.
use anyhow::{Context, Result, ensure};
use axum::http::{HeaderValue, Method, header};
use std::{
    env,
    path::Path,
    sync::{Arc, Mutex},
};
use tokenling_server::{
    brain::Brain,
    host::{Host, Worker},
    store::Store,
};
use tower_http::{cors::CorsLayer, set_header::SetResponseHeaderLayer};
fn required(name: &str) -> Result<String> {
    env::var(name).with_context(|| format!("set {name}"))
}
#[tokio::main]
async fn main() -> Result<()> {
    let listen: std::net::SocketAddr = required("SUBTOKENLING_LISTEN")?.parse()?;
    let db = required("SUBTOKENLING_DB")?;
    let model_dir = required("SUBTOKENLING_MODEL_DIR")?;
    let model_id = required("SUBTOKENLING_MODEL_ID")?;
    let dtype = env::var("SUBTOKENLING_MODEL_DTYPE")
        .unwrap_or_else(|_| "q4".into())
        .parse()?;
    let origins: Vec<HeaderValue> = required("SUBTOKENLING_ALLOWED_ORIGIN")?
        .split(',')
        .map(|s| s.trim().parse())
        .collect::<std::result::Result<_, _>>()?;
    ensure!(
        !origins.is_empty() && origins.iter().all(|s| s.as_bytes() != b"*"),
        "configure explicit allowed browser origins"
    );
    let threads: usize = env::var("SUBTOKENLING_THREADS")
        .unwrap_or_else(|_| "4".into())
        .parse()?;
    ensure!(threads > 0, "SUBTOKENLING_THREADS must be positive");
    let max_pets: usize = env::var("SUBTOKENLING_MAX_PETS")
        .unwrap_or_else(|_| "64".into())
        .parse()?;
    ensure!(max_pets > 0, "SUBTOKENLING_MAX_PETS must be positive");
    let store = Store::open(Path::new(&db))?;
    let foreign = store.foreign_brains(&model_id)?;
    ensure!(
        foreign == 0,
        "{foreign} hosted pet(s) run a brain other than SUBTOKENLING_MODEL_ID; keep their model, \
         or have owners bring them local before changing it (no pet was changed)"
    );
    ort::init_from(required("ORT_DYLIB_PATH")?)?
        .with_name("tokenling")
        .commit();
    let cpu = Worker::new(
        "cpu",
        Brain::load(Path::new(&model_dir), dtype, threads, None)?,
    );
    let gpu = if let Ok(device) = env::var("SUBTOKENLING_GPU_DEVICE") {
        match Brain::load(Path::new(&model_dir), dtype, 1, Some(device.parse()?)) {
            Ok(brain) => Some(Worker::new("gpu", brain)),
            Err(error) => {
                eprintln!("optional CUDA worker unavailable: {error}");
                None
            }
        }
    } else {
        None
    };
    let host = Arc::new(Host {
        db: Mutex::new(store),
        model_id,
        features: serde_json::from_str(
            &env::var("SUBTOKENLING_FEATURES").unwrap_or_else(|_| "[]".into()),
        )?,
        max_pets,
        weather_seed: env::var("SUBTOKENLING_WEATHER_SEED").unwrap_or_default(),
        cpu,
        gpu,
        jobs: Mutex::new(std::collections::HashMap::new()),
        notify: tokio::sync::Notify::new(),
        queue: Mutex::new(Default::default()),
    });
    let router = host
        .clone()
        .router()
        .layer(
            CorsLayer::new()
                .allow_origin(origins)
                .allow_methods([Method::GET, Method::POST, Method::DELETE])
                .allow_headers([header::AUTHORIZATION, header::CONTENT_TYPE]),
        )
        .layer(SetResponseHeaderLayer::if_not_present(
            header::CACHE_CONTROL,
            HeaderValue::from_static("no-store"),
        ))
        .layer(SetResponseHeaderLayer::if_not_present(
            header::HeaderName::from_static("cross-origin-resource-policy"),
            HeaderValue::from_static("cross-origin"),
        ));
    let listener = tokio::net::TcpListener::bind(listen).await?;
    let scheduler = tokio::spawn(host.run());
    eprintln!("subtokenling host listening on {listen}");
    axum::serve(listener, router)
        .with_graceful_shutdown(async {
            let _ = tokio::signal::ctrl_c().await;
        })
        .await?;
    scheduler.abort();
    Ok(())
}