//! Inbound federation handler: implements `FederationService` Cap'n Proto interface.
//!
//! Delegates all operations to the local [`Store`], acting as a trusted relay
//! from authenticated peer servers.

use std::sync::Arc;

use capnp::capability::Promise;
use quicproquo_proto::federation_capnp::federation_service;
use tokio::sync::Notify;
use dashmap::DashMap;

use crate::storage::Store;

/// Inbound federation RPC handler.
pub struct FederationServiceImpl {
    pub store: Arc<dyn Store>,
    pub waiters: Arc<DashMap<Vec<u8>, Arc<Notify>>>,
    pub local_domain: String,
}

impl federation_service::Server for FederationServiceImpl {
    fn relay_enqueue(
        &mut self,
        params: federation_service::RelayEnqueueParams,
        mut results: federation_service::RelayEnqueueResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        let recipient_key = match p.get_recipient_key() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad recipient_key: {e}"))),
        };
        let payload = match p.get_payload() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad payload: {e}"))),
        };
        let channel_id = p.get_channel_id().unwrap_or_default().to_vec();

        if let Ok(a) = p.get_auth() {
            if let Ok(origin) = a.get_origin() {
                let origin = origin.to_str().unwrap_or("?");
                tracing::debug!(origin = origin, "federation relay_enqueue");
            }
        }

        if recipient_key.len() != 32 {
            return Promise::err(capnp::Error::failed("recipient_key must be 32 bytes".into()));
        }
        if payload.is_empty() {
            return Promise::err(capnp::Error::failed("payload must not be empty".into()));
        }

        let seq = match self.store.enqueue(&recipient_key, &channel_id, payload, None) {
            Ok(s) => s,
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        };

        results.get().set_seq(seq);

        // Wake any waiting fetchWait clients.
        if let Some(waiter) = self.waiters.get(&recipient_key) {
            waiter.notify_waiters();
        }

        tracing::info!(
            recipient_prefix = %hex::encode(&recipient_key[..4]),
            seq = seq,
            "federation: relayed enqueue"
        );

        Promise::ok(())
    }

    fn relay_batch_enqueue(
        &mut self,
        params: federation_service::RelayBatchEnqueueParams,
        mut results: federation_service::RelayBatchEnqueueResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        let recipient_keys = match p.get_recipient_keys() {
            Ok(v) => v,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad recipient_keys: {e}"))),
        };
        let payload = match p.get_payload() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad payload: {e}"))),
        };
        let channel_id = p.get_channel_id().unwrap_or_default().to_vec();

        let mut seqs = Vec::with_capacity(recipient_keys.len() as usize);
        for i in 0..recipient_keys.len() {
            let rk = match recipient_keys.get(i) {
                Ok(v) => v.to_vec(),
                Err(e) => return Promise::err(capnp::Error::failed(format!("bad key[{i}]: {e}"))),
            };
            if rk.len() != 32 {
                return Promise::err(capnp::Error::failed(
                    format!("recipient_key[{i}] must be 32 bytes"),
                ));
            }
            let seq = match self.store.enqueue(&rk, &channel_id, payload.clone(), None) {
                Ok(s) => s,
                Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
            };
            seqs.push(seq);
            if let Some(waiter) = self.waiters.get(&rk) {
                waiter.notify_waiters();
            }
        }

        let mut list = results.get().init_seqs(seqs.len() as u32);
        for (i, seq) in seqs.iter().enumerate() {
            list.set(i as u32, *seq);
        }

        tracing::info!(
            recipient_count = recipient_keys.len(),
            "federation: relayed batch_enqueue"
        );

        Promise::ok(())
    }

    fn proxy_fetch_key_package(
        &mut self,
        params: federation_service::ProxyFetchKeyPackageParams,
        mut results: federation_service::ProxyFetchKeyPackageResults,
    ) -> Promise<(), capnp::Error> {
        let identity_key = match params.get().and_then(|p| p.get_identity_key()) {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        match self.store.fetch_key_package(&identity_key) {
            Ok(Some(pkg)) => results.get().set_package(&pkg),
            Ok(None) => results.get().set_package(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn proxy_fetch_hybrid_key(
        &mut self,
        params: federation_service::ProxyFetchHybridKeyParams,
        mut results: federation_service::ProxyFetchHybridKeyResults,
    ) -> Promise<(), capnp::Error> {
        let identity_key = match params.get().and_then(|p| p.get_identity_key()) {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        match self.store.fetch_hybrid_key(&identity_key) {
            Ok(Some(pk)) => results.get().set_hybrid_public_key(&pk),
            Ok(None) => results.get().set_hybrid_public_key(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn proxy_resolve_user(
        &mut self,
        params: federation_service::ProxyResolveUserParams,
        mut results: federation_service::ProxyResolveUserResults,
    ) -> Promise<(), capnp::Error> {
        let username = match params.get().and_then(|p| p.get_username()) {
            Ok(u) => match u.to_str() {
                Ok(s) => s.to_string(),
                Err(e) => return Promise::err(capnp::Error::failed(format!("bad utf-8: {e}"))),
            },
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        match self.store.get_user_identity_key(&username) {
            Ok(Some(key)) => results.get().set_identity_key(&key),
            Ok(None) => results.get().set_identity_key(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn federation_health(
        &mut self,
        _params: federation_service::FederationHealthParams,
        mut results: federation_service::FederationHealthResults,
    ) -> Promise<(), capnp::Error> {
        results.get().set_status("ok");
        results.get().set_server_domain(&self.local_domain);
        Promise::ok(())
    }
}
