Lifecycle

use std::collections::HashMap;
use std::hash::DefaultHasher;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};

use kanau::{JsonMessageDe, JsonMessageSer};
use serde::{Deserialize, Serialize};
use wakuwaku::deps::kanau::processor::Processor;
use wakuwaku::deps::tonic::{self, Request, Response, Status};
use wakuwaku::iggy::error::HandleError;
use wakuwaku::iggy::{
    Event, EventAtomicOrdering, EventBody, EventHandler, EventTypeTag, IggySuperPartitionName,
    NonAssociative, PartitionKey,
};
use wakuwaku::{CronJob, CronJobInterval, CronTick, ServiceCreation};

use crate::proto::accounts_server::Accounts;
use crate::proto::{Balance, GetBalanceRequest};

// ---- Shared dependencies: cheap to clone, every clone shares the same state.

/// Account balances. Stands in for a database handle.
#[derive(Clone, Default)]
pub struct AccountStore {
    balances: Arc<Mutex<HashMap<u64, i64>>>,
}

impl AccountStore {
    pub fn deposit(&self, account: u64, cents: i64) {
        *self.balances.lock().unwrap().entry(account).or_default() += cents;
    }

    pub fn balance(&self, account: u64) -> Option<i64> {
        self.balances.lock().unwrap().get(&account).copied()
    }
}

/// Login sessions and when they expire.
#[derive(Clone, Default)]
pub struct SessionStore {
    expires_at: Arc<Mutex<HashMap<String, SystemTime>>>,
}

impl SessionStore {
    pub fn insert(&self, token: &str, expires_at: SystemTime) {
        self.expires_at.lock().unwrap().insert(token.to_owned(), expires_at);
    }

    pub fn contains(&self, token: &str) -> bool {
        self.expires_at.lock().unwrap().contains_key(token)
    }

    /// Remove every session that expired before `now`; return how many.
    pub fn remove_expired(&self, now: SystemTime) -> usize {
        let mut sessions = self.expires_at.lock().unwrap();
        let before = sessions.len();
        sessions.retain(|_, expires_at| *expires_at > now);
        before - sessions.len()
    }
}

/// Outgoing mail. Stands in for an SMTP client.
#[derive(Clone, Default)]
pub struct Mailer {
    sent: Arc<Mutex<Vec<String>>>,
}

impl Mailer {
    pub fn send(&self, to: &str) {
        self.sent.lock().unwrap().push(to.to_owned());
    }

    pub fn sent(&self) -> Vec<String> {
        self.sent.lock().unwrap().clone()
    }
}

// ---- Cron job: expire sessions every minute.

pub struct ExpireSessions {
    sessions: SessionStore,
}

impl ServiceCreation for ExpireSessions {
    type Dep = SessionStore;
    fn create_service(sessions: SessionStore) -> Self {
        Self { sessions }
    }
}

impl Processor<CronTick> for ExpireSessions {
    type Output = ();
    type Error = std::convert::Infallible;
    async fn process(&self, tick: CronTick) -> Result<(), Self::Error> {
        let removed = self.sessions.remove_expired(tick.fired_at);
        tracing::info!(removed, "expired sessions");
        Ok(())
    }
}

impl CronJob for ExpireSessions {
    const INTERVAL: CronJobInterval = CronJobInterval::AtLeast(Duration::from_secs(60));
}

// ---- Event handler: welcome mail for every new user.

#[derive(Hash)]
pub struct UserId(pub u64);

impl PartitionKey for UserId {
    type PartitionMerge = DefaultHasher;
    const SUPER_PARTITION_NAME: IggySuperPartitionName = IggySuperPartitionName("users");
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonMessageSer, JsonMessageDe)]
pub struct UserRegistered {
    pub user_id: u64,
    pub email: String,
}

impl EventBody for UserRegistered {}

impl Event for UserRegistered {
    const TYPE_TAG: EventTypeTag = EventTypeTag::new(1);
    type Key = UserId;
    type Algebra = NonAssociative;
    fn key(&self) -> UserId {
        UserId(self.user_id)
    }
    fn atomic_ordering(&self) -> EventAtomicOrdering {
        EventAtomicOrdering::Relaxed
    }
}

pub struct SendWelcomeMail {
    mailer: Mailer,
}

impl ServiceCreation for SendWelcomeMail {
    type Dep = Mailer;
    fn create_service(mailer: Mailer) -> Self {
        Self { mailer }
    }
}

impl EventHandler<UserRegistered> for SendWelcomeMail {
    const MAX_BATCH: NonZeroUsize = NonZeroUsize::new(32).unwrap();
    async fn handle(&self, run: &[UserRegistered]) -> Result<(), HandleError> {
        for event in run {
            self.mailer.send(&event.email);
        }
        Ok(())
    }
}

// ---- gRPC service: read balances.

pub struct AccountsService {
    store: AccountStore,
}

impl ServiceCreation for AccountsService {
    type Dep = AccountStore;
    fn create_service(store: AccountStore) -> Self {
        Self { store }
    }
}

#[tonic::async_trait]
impl Accounts for AccountsService {
    async fn get_balance(
        &self,
        request: Request<GetBalanceRequest>,
    ) -> Result<Response<Balance>, Status> {
        let account_id = request.into_inner().account_id;
        let cents = self
            .store
            .balance(account_id)
            .ok_or_else(|| Status::not_found("no such account"))?;
        Ok(Response::new(Balance { account_id, cents }))
    }
}

mod components;

pub mod proto {
    wakuwaku::deps::tonic::include_proto!("accounts.v1");
}

use std::sync::Arc;

use tracing_subscriber::EnvFilter;
use wakuwaku::ServiceBuilder;
use wakuwaku::deps::iggy::prelude::{Client, IggyClient};
use wakuwaku::deps::tonic::service::Routes;
use wakuwaku::deps::tonic::transport::Server;
use wakuwaku::deps::tonic::transport::server::TcpIncoming;
use wakuwaku::iggy::ConsumerConfig;

use components::{
    AccountStore, AccountsService, ExpireSessions, Mailer, SendWelcomeMail, SessionStore,
};
use proto::accounts_server::AccountsServer;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 0. Logging first, so every later step can report.
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()))
        .init();

    // 1. Connect to infrastructure and bind the port. Any error here stops
    //    startup before a single component runs.
    let iggy = Arc::new(IggyClient::from_connection_string(
        "iggy://iggy:iggy@localhost:8090",
    )?);
    iggy.connect().await?;
    let incoming = TcpIncoming::bind("[::]:50051".parse()?)?;

    // 2. Register the shared dependencies.
    let services = ServiceBuilder::new(AccountStore::default())
        .push(SessionStore::default())
        .push(Mailer::default());

    // 3. Event consumers. `start` validates the topics and can fail.
    let (services, consumers) = services
        .iggy_consumer::<SendWelcomeMail, _, _>()
        .start(iggy.clone(), ConsumerConfig::new("shop", "accounts", vec![0]))
        .await?;

    // 4. gRPC servers. Built now, served in step 6.
    let (services, add_grpc) = services
        .grpc_service::<AccountsServer<_>, AccountsService, _>(AccountsServer::new)
        .into_router_fn();

    // 5. Cron jobs. The first run of every job happens right away.
    let (_services, cron) = services.cron_job::<ExpireSessions, _>().start();

    // 6. Serve until SIGTERM / Ctrl-C. In-flight requests finish first.
    let router = add_grpc(Server::builder().add_routes(Routes::default()));
    let served = router
        .serve_with_incoming_shutdown(incoming, shutdown_signal())
        .await;

    // 7. Drain the rest, then report how serving ended.
    consumers.shutdown().await;
    cron.shutdown().await;
    served?;
    tracing::info!("stopped");
    Ok(())
}

async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c().await.expect("install Ctrl-C handler");
    };
    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("install SIGTERM handler")
            .recv()
            .await;
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        () = ctrl_c => {}
        () = terminate => {}
    }
    tracing::info!("shutdown signal received");
}

// Give the background work at most 20 seconds; past that, cancel it.
if tokio::time::timeout(Duration::from_secs(20), consumers.shutdown()).await.is_err() {
    tracing::warn!("consumers did not drain in time");
}