/*
 * asklyphe-common nats/comms.rs
 * - communication functions for nats
 *
 * Copyright (C) 2025 Real Microsoft, LLC
 *
 * This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
*/

use crate::nats::authservice::{AUTH_REPLY_SERVICE, AUTH_SERVICE};
use crate::nats::bingservice::{BingServiceResponse, BING_REPLY_SERVICE, BING_SERVICE, GOOGLE_REPLY_SERVICE, GOOGLE_SERVICE};
use crate::nats::searchservice::{SearchSrvcResponse, SEARCH_REPLY_SERVICE, SEARCH_SERVICE};
use crate::nats::{authservice, bingservice, searchservice};
use async_nats::jetstream;
use async_nats::jetstream::consumer::PullConsumer;
use async_nats::jetstream::stream::RetentionPolicy;
use futures::StreamExt;
use log::{debug, error, warn};
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::future::Future;
use std::hash::{DefaultHasher, Hash, Hasher};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use ulid::Ulid;

#[derive(Serialize, Deserialize, Clone)]
pub enum Query {
    SearchService(searchservice::SearchSrvcQuery),
    BingService(bingservice::BingServiceQuery),
    GoogleService(bingservice::BingServiceQuery),
    AuthService(authservice::AuthServiceQuery),
}

impl PartialEq<Self> for Query {
    fn eq(&self, other: &Self) -> bool {
        let a = rmp_serde::to_vec(self).unwrap();
        let b = rmp_serde::to_vec(other).unwrap();
        a == b
    }
}

impl Eq for Query {}
impl PartialOrd<Self> for Query {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        let a = rmp_serde::to_vec(self).unwrap();
        let b = rmp_serde::to_vec(other).unwrap();
        let mut hash = DefaultHasher::new();
        a.hash(&mut hash);
        let a = hash.finish();
        let mut hash = DefaultHasher::new();
        b.hash(&mut hash);
        let b = hash.finish();
        a.partial_cmp(&b)
    }
}

impl Ord for Query {
    fn cmp(&self, other: &Self) -> Ordering {
        let a = rmp_serde::to_vec(self).unwrap();
        let b = rmp_serde::to_vec(other).unwrap();
        let mut hash = DefaultHasher::new();
        a.hash(&mut hash);
        let a = hash.finish();
        let mut hash = DefaultHasher::new();
        b.hash(&mut hash);
        let b = hash.finish();
        a.cmp(&b)
    }
}

pub enum Service {
    SearchService,
    BingService,
    GoogleService, // mirror of bingservice but with different name
    AuthService,
}

#[derive(Clone)]
pub enum ServiceResponse {
    SearchService(searchservice::SearchSrvcResponse),
    BingService(bingservice::BingServiceResponse),
    AuthService(authservice::AuthServiceResponse),
}

#[derive(Copy, Clone, Debug)]
pub enum ServiceError {
    /// may indicate that the given nats server is down
    NatsSubscriptionFailure,
    /// may indicate that the given nats server is out of memory, or overloaded
    NatsPublishFailure,
    /// could mean multiple things, but treat this as nats being down
    NatsRecvMessageFailure,
    /// usually this one happens if we try to get a service that doesn't exist, but i haven't
    /// experienced it lately in production
    NatsGetServiceError,
    /// never experienced in production, unsure when/if this could occur
    SerialiseError,
    /// indicates something wrong with a service / backend, check logs
    DeserialiseError,
}

pub fn ulid() -> String {
    Ulid::new().to_string()
}

static SERVICE_CACHE: RwLock<BTreeMap<Query, ServiceResponse>> = RwLock::const_new(BTreeMap::new());
pub const CACHE_PERIOD_MINUTES: u64 = 120;

// todo: add timeouts
/// use_cache will enable deterministic caching, where when given the same input, the same output *should* happen
/// only use for search queries and such, not anything authentication related
pub async fn query_service(
    mut query: Query,
    nats: &jetstream::Context,
    use_cache: bool,
) -> Result<ServiceResponse, ServiceError> {
    {
        if use_cache {
            let cache = SERVICE_CACHE.read().await;
            let mut query = query.clone();
            match &mut query {
                Query::SearchService(s) => s.replyto = "".to_string(),
                Query::BingService(s) => s.replyto = "".to_string(),
                Query::GoogleService(s) => s.replyto = "".to_string(),
                Query::AuthService(s) => s.replyto = "".to_string(),
            }
            if let Some(r) = cache.get(&query).cloned() {
                debug!("cache hit");
                return Ok(r);
            } else {
                debug!("cache miss");
            }
        }
    }

    let ulid = ulid();
    let service = match &query {
        Query::SearchService(_) => SEARCH_SERVICE,
        Query::AuthService(_) => AUTH_SERVICE,
        Query::BingService(_) => BING_SERVICE,
        Query::GoogleService(_) => GOOGLE_SERVICE,
    };
    let replyservice = match &query {
        Query::SearchService(_) => SEARCH_REPLY_SERVICE,
        Query::AuthService(_) => AUTH_REPLY_SERVICE,
        Query::BingService(_) => BING_REPLY_SERVICE,
        Query::GoogleService(_) => GOOGLE_REPLY_SERVICE,
    };

    let stream = nats
        .get_or_create_stream(jetstream::stream::Config {
            name: replyservice.to_string(),
            subjects: vec![format!("{}.*", replyservice)],
            retention: RetentionPolicy::Limits,
            max_messages_per_subject: 1,
            max_age: Duration::from_secs(60),
            ..Default::default()
        })
        .await
        .map_err(|_| ServiceError::NatsGetServiceError)?;

    let consumer_name = format!("{}-{}", replyservice, ulid);

    // todo: kinda dumb way to do things, maybe change?
    let reply_to = format!("{replyservice}.{}", ulid);
    match &mut query {
        Query::SearchService(q) => q.replyto = reply_to.clone(),
        Query::AuthService(q) => q.replyto = reply_to.clone(),
        Query::BingService(q) => q.replyto = reply_to.clone(),
        Query::GoogleService(q) => q.replyto = reply_to.clone(),
    }

    let consumer: PullConsumer = stream
        .get_or_create_consumer(
            &consumer_name,
            jetstream::consumer::pull::Config {
                filter_subject: reply_to.clone(),
                inactive_threshold: Duration::from_secs(65), // todo: choose a less arbitrary value?
                durable_name: None,
                ..Default::default()
            },
        )
        .await
        .map_err(|_| ServiceError::NatsSubscriptionFailure)?;

    // todo! log these errors if they ever occur
    let buf = match &query {
        Query::SearchService(q) => {
            rmp_serde::encode::to_vec(q).map_err(|_| ServiceError::SerialiseError)?
        }
        Query::AuthService(q) => {
            rmp_serde::encode::to_vec(q).map_err(|_| ServiceError::SerialiseError)?
        }
        Query::BingService(q) => {
            rmp_serde::encode::to_vec(q).map_err(|_| ServiceError::SerialiseError)?
        }
        Query::GoogleService(q) => {
            rmp_serde::encode::to_vec(q).map_err(|_| ServiceError::SerialiseError)?
        }
    };

    if let Err(e) = nats.publish(service, buf.into()).await {
        error!("error occurred while trying to publish to nats!! {e}");
        return Err(ServiceError::NatsPublishFailure);
    }

    // todo: log these errors?
    let mut messages = consumer
        .messages()
        .await
        .map_err(|_| ServiceError::NatsSubscriptionFailure)?;

    if let Some(Ok(message)) = messages.next().await {
        if let Err(e) = message.ack().await {
            warn!("failure while acking message! {e}");
        }

        let result = match &query {
            Query::SearchService(_) => ServiceResponse::SearchService(
                rmp_serde::from_slice::<searchservice::SearchSrvcResponse>(
                    message.payload.as_ref(),
                )
                .map_err(|_| ServiceError::DeserialiseError)?,
            ),
            Query::AuthService(_) => ServiceResponse::AuthService(
                rmp_serde::from_slice::<authservice::AuthServiceResponse>(message.payload.as_ref())
                    .map_err(|_| ServiceError::DeserialiseError)?,
            ),
            Query::BingService(_) => ServiceResponse::BingService(
                rmp_serde::from_slice::<bingservice::BingServiceResponse>(message.payload.as_ref())
                    .map_err(|_| ServiceError::DeserialiseError)?,
            ),
            Query::GoogleService(_) => ServiceResponse::BingService(
                rmp_serde::from_slice::<bingservice::BingServiceResponse>(message.payload.as_ref())
                    .map_err(|_| ServiceError::DeserialiseError)?,
            ),
        };
        
        if use_cache {
            // only write cache if this isn't an error
            let mut error = false;
            match &result {
                ServiceResponse::SearchService(s) => match s {
                    SearchSrvcResponse::InvalidRequest => { error = true; }
                    SearchSrvcResponse::OtherError(_) => { error = true; }
                    _ => {}
                }
                ServiceResponse::BingService(s) => match s {
                    BingServiceResponse::InvalidRequest => { error = true; }
                    BingServiceResponse::OtherError(_) => { error = true; }
                    _ => {}
                }
                ServiceResponse::AuthService(_) => { error = true; } // fallback to make sure we absolutely don't cache authservice stuff
            }
            if !error {
                let mut cache = SERVICE_CACHE.write().await;
                let mut query = query.clone();
                match &mut query {
                    Query::SearchService(s) => s.replyto = "".to_string(),
                    Query::BingService(s) => s.replyto = "".to_string(),
                    Query::GoogleService(s) => s.replyto = "".to_string(),
                    Query::AuthService(s) => s.replyto = "".to_string(),
                }
                cache.insert(query.clone(), result.clone());
                let query_clone = query.clone();
                tokio::spawn(async move {
                    tokio::time::sleep(Duration::from_secs(60 * CACHE_PERIOD_MINUTES)).await;
                    let mut cache = SERVICE_CACHE.write().await;
                    cache.remove(&query_clone);
                });
            }
        }
        
        Ok(result)
    } else {
        error!("couldn't receive message from service {} !!!!", service);
        Err(ServiceError::NatsRecvMessageFailure)
    }
}

pub async fn subscribe_service<
    F: Fn(Query) -> Fut + Sync + Send + 'static,
    Fut: Send + Future<Output = ServiceResponse>,
>(
    service: Service,
    nats: &jetstream::Context,
    process: Arc<F>,
) -> Result<(), ServiceError> {
    let service_name = match &service {
        Service::SearchService => SEARCH_SERVICE,
        Service::AuthService => AUTH_SERVICE,
        Service::BingService => BING_SERVICE,
        Service::GoogleService => GOOGLE_SERVICE,
    };

    let consumer: PullConsumer = nats
        .get_or_create_stream(jetstream::stream::Config {
            name: service_name.to_string(),
            subjects: vec![service_name.to_string()],
            retention: RetentionPolicy::WorkQueue,
            ..Default::default()
        })
        .await
        .map_err(|_| ServiceError::NatsGetServiceError)?
        .get_or_create_consumer(
            &format!("{}er", service_name),
            jetstream::consumer::pull::Config {
                durable_name: Some(format!("{}er", service_name)),
                filter_subject: service_name.to_string(),
                ..Default::default()
            },
        )
        .await
        .map_err(|_| ServiceError::NatsGetServiceError)?;

    let mut messages = consumer
        .messages()
        .await
        .map_err(|_| ServiceError::NatsSubscriptionFailure)?;

    while let Some(message) = messages.next().await {
        if let Err(e) = message {
            warn!("error when recv js message! {e}");
            continue;
        }
        let message = message.unwrap();

        if let Err(e) = message.ack().await {
            warn!("failure while acking message! {e}");
        }

        let message = match &service {
            Service::SearchService => {
                let res = rmp_serde::from_slice::<searchservice::SearchSrvcQuery>(
                    message.payload.as_ref(),
                );
                if let Err(e) = res {
                    warn!("invalid message sent via js!!! {e}");
                    continue;
                }
                Query::SearchService(res.unwrap())
            }
            Service::AuthService => {
                let res = rmp_serde::from_slice::<authservice::AuthServiceQuery>(
                    message.payload.as_ref(),
                );
                if let Err(e) = res {
                    warn!("invalid message sent via js!!! {e}");
                    continue;
                }
                Query::AuthService(res.unwrap())
            }
            Service::BingService => {
                let res = rmp_serde::from_slice::<bingservice::BingServiceQuery>(
                    message.payload.as_ref(),
                );
                if let Err(e) = res {
                    warn!("invalid message sent via js!!! {e}");
                    continue;
                }
                Query::BingService(res.unwrap())
            }
            Service::GoogleService => {
                let res = rmp_serde::from_slice::<bingservice::BingServiceQuery>(
                    message.payload.as_ref(),
                );
                if let Err(e) = res {
                    warn!("invalid message sent via js!!! {e}");
                    continue;
                }
                Query::GoogleService(res.unwrap())
            }
        };

        let replyto = match &message {
            Query::SearchService(m) => m.replyto.clone(),
            Query::AuthService(m) => m.replyto.clone(),
            Query::BingService(m) => m.replyto.clone(),
            Query::GoogleService(m) => m.replyto.clone(),
        };

        let nats = nats.clone();
        let process = process.clone();
        tokio::spawn(async move {
            let response = process(message).await;
            let buf = match &response {
                ServiceResponse::SearchService(r) => {
                    let res = rmp_serde::encode::to_vec(r);
                    if let Err(e) = res {
                        warn!("unable to serialise response!!! {e}");
                        return;
                    }
                    res.unwrap()
                }
                ServiceResponse::AuthService(r) => {
                    let res = rmp_serde::encode::to_vec(r);
                    if let Err(e) = res {
                        warn!("unable to serialise response!!! {e}");
                        return;
                    }
                    res.unwrap()
                }
                ServiceResponse::BingService(r) => {
                    let res = rmp_serde::encode::to_vec(r);
                    if let Err(e) = res {
                        warn!("unable to serialise response!!! {e}");
                        return;
                    }
                    res.unwrap()
                }
            };

            if let Err(e) = nats.publish(replyto, buf.into()).await {
                error!("error occurred while trying to publish to nats!! {e}");
            }
        });
    }

    Ok(())
}
