//! Outbound federation client: connects to peer servers to relay messages.
//!
//! Uses a lazy connection pool (DashMap) to reuse QUIC connections to known peers.

#![allow(dead_code)] // federation not yet wired up

use std::collections::HashMap;
use std::net::SocketAddr;

use anyhow::Context;
use dashmap::DashMap;
use quinn::Endpoint;

use crate::config::EffectiveFederationConfig;

/// Outbound federation client for relaying to peer servers.
pub struct FederationClient {
    /// Peer domain → address mapping from config.
    peer_addresses: HashMap<String, SocketAddr>,
    /// Lazy QUIC connection pool: domain → active Connection.
    connections: DashMap<String, quinn::Connection>,
    /// Local QUIC endpoint (shared for all outbound federation connections).
    endpoint: Endpoint,
    /// Local domain (for the FederationAuth.origin field).
    local_domain: String,
}

impl FederationClient {
    /// Create a new federation client from config.
    ///
    /// The `endpoint` should be configured with mTLS client credentials.
    pub fn new(
        config: &EffectiveFederationConfig,
        endpoint: Endpoint,
    ) -> anyhow::Result<Self> {
        let mut peer_addresses = HashMap::new();
        for peer in &config.peers {
            let addr: SocketAddr = peer.address.parse().with_context(|| {
                format!("parse federation peer address '{}' for '{}'", peer.address, peer.domain)
            })?;
            peer_addresses.insert(peer.domain.clone(), addr);
        }

        Ok(Self {
            peer_addresses,
            connections: DashMap::new(),
            endpoint,
            local_domain: config.domain.clone(),
        })
    }

    /// Check if we have a configured peer for the given domain.
    pub fn has_peer(&self, domain: &str) -> bool {
        self.peer_addresses.contains_key(domain)
    }

    /// List all configured peer domains.
    pub fn peer_domains(&self) -> Vec<String> {
        self.peer_addresses.keys().cloned().collect()
    }

    /// Get the local domain.
    pub fn local_domain(&self) -> &str {
        &self.local_domain
    }

    /// Relay a single enqueue to a remote peer. Returns the seq assigned by the remote server.
    pub async fn relay_enqueue(
        &self,
        domain: &str,
        recipient_key: &[u8],
        payload: &[u8],
        channel_id: &[u8],
    ) -> anyhow::Result<u64> {
        let conn = self.get_or_connect(domain).await?;
        let (send, recv) = conn.open_bi().await.context("open bi stream to peer")?;

        let (reader, writer) = (
            tokio_util::compat::TokioAsyncReadCompatExt::compat(recv),
            tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(send),
        );

        let rpc_network = capnp_rpc::twoparty::VatNetwork::new(
            reader,
            writer,
            capnp_rpc::rpc_twoparty_capnp::Side::Client,
            Default::default(),
        );

        let mut rpc_system = capnp_rpc::RpcSystem::new(Box::new(rpc_network), None);
        let client: quicproquo_proto::federation_capnp::federation_service::Client =
            rpc_system.bootstrap(capnp_rpc::rpc_twoparty_capnp::Side::Server);

        tokio::task::spawn_local(rpc_system);

        let mut req = client.relay_enqueue_request();
        {
            let mut builder = req.get();
            builder.set_recipient_key(recipient_key);
            builder.set_payload(payload);
            builder.set_channel_id(channel_id);
            builder.set_version(1);
            let mut auth = builder.init_auth();
            auth.set_origin(&self.local_domain);
        }

        let response = req.send().promise.await
            .map_err(|e| anyhow::anyhow!("federation relay_enqueue failed: {e}"))?;
        let seq = response.get()
            .map_err(|e| anyhow::anyhow!("read relay_enqueue response: {e}"))?
            .get_seq();

        Ok(seq)
    }

    /// Proxy a key package fetch to a remote peer.
    pub async fn proxy_fetch_key_package(
        &self,
        domain: &str,
        identity_key: &[u8],
    ) -> anyhow::Result<Option<Vec<u8>>> {
        let conn = self.get_or_connect(domain).await?;
        let (send, recv) = conn.open_bi().await.context("open bi stream to peer")?;

        let (reader, writer) = (
            tokio_util::compat::TokioAsyncReadCompatExt::compat(recv),
            tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(send),
        );

        let rpc_network = capnp_rpc::twoparty::VatNetwork::new(
            reader,
            writer,
            capnp_rpc::rpc_twoparty_capnp::Side::Client,
            Default::default(),
        );

        let mut rpc_system = capnp_rpc::RpcSystem::new(Box::new(rpc_network), None);
        let client: quicproquo_proto::federation_capnp::federation_service::Client =
            rpc_system.bootstrap(capnp_rpc::rpc_twoparty_capnp::Side::Server);

        tokio::task::spawn_local(rpc_system);

        let mut req = client.proxy_fetch_key_package_request();
        {
            let mut builder = req.get();
            builder.set_identity_key(identity_key);
            let mut auth = builder.init_auth();
            auth.set_origin(&self.local_domain);
        }

        let response = req.send().promise.await
            .map_err(|e| anyhow::anyhow!("federation proxy_fetch_key_package failed: {e}"))?;
        let pkg = response.get()
            .map_err(|e| anyhow::anyhow!("read proxy_fetch_key_package response: {e}"))?
            .get_package()
            .map_err(|e| anyhow::anyhow!("get package: {e}"))?;

        if pkg.is_empty() {
            Ok(None)
        } else {
            Ok(Some(pkg.to_vec()))
        }
    }

    /// Proxy a hybrid key fetch to a remote peer.
    pub async fn proxy_fetch_hybrid_key(
        &self,
        domain: &str,
        identity_key: &[u8],
    ) -> anyhow::Result<Option<Vec<u8>>> {
        let conn = self.get_or_connect(domain).await?;
        let (send, recv) = conn.open_bi().await.context("open bi stream to peer")?;

        let (reader, writer) = (
            tokio_util::compat::TokioAsyncReadCompatExt::compat(recv),
            tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(send),
        );

        let rpc_network = capnp_rpc::twoparty::VatNetwork::new(
            reader,
            writer,
            capnp_rpc::rpc_twoparty_capnp::Side::Client,
            Default::default(),
        );

        let mut rpc_system = capnp_rpc::RpcSystem::new(Box::new(rpc_network), None);
        let client: quicproquo_proto::federation_capnp::federation_service::Client =
            rpc_system.bootstrap(capnp_rpc::rpc_twoparty_capnp::Side::Server);

        tokio::task::spawn_local(rpc_system);

        let mut req = client.proxy_fetch_hybrid_key_request();
        {
            let mut builder = req.get();
            builder.set_identity_key(identity_key);
            let mut auth = builder.init_auth();
            auth.set_origin(&self.local_domain);
        }

        let response = req.send().promise.await
            .map_err(|e| anyhow::anyhow!("federation proxy_fetch_hybrid_key failed: {e}"))?;
        let pk = response.get()
            .map_err(|e| anyhow::anyhow!("read proxy_fetch_hybrid_key response: {e}"))?
            .get_hybrid_public_key()
            .map_err(|e| anyhow::anyhow!("get hybrid_public_key: {e}"))?;

        if pk.is_empty() {
            Ok(None)
        } else {
            Ok(Some(pk.to_vec()))
        }
    }

    /// Proxy a user resolution to a remote peer.
    pub async fn proxy_resolve_user(
        &self,
        domain: &str,
        username: &str,
    ) -> anyhow::Result<Option<Vec<u8>>> {
        let conn = self.get_or_connect(domain).await?;
        let (send, recv) = conn.open_bi().await.context("open bi stream to peer")?;

        let (reader, writer) = (
            tokio_util::compat::TokioAsyncReadCompatExt::compat(recv),
            tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(send),
        );

        let rpc_network = capnp_rpc::twoparty::VatNetwork::new(
            reader,
            writer,
            capnp_rpc::rpc_twoparty_capnp::Side::Client,
            Default::default(),
        );

        let mut rpc_system = capnp_rpc::RpcSystem::new(Box::new(rpc_network), None);
        let client: quicproquo_proto::federation_capnp::federation_service::Client =
            rpc_system.bootstrap(capnp_rpc::rpc_twoparty_capnp::Side::Server);

        tokio::task::spawn_local(rpc_system);

        let mut req = client.proxy_resolve_user_request();
        {
            let mut builder = req.get();
            builder.set_username(username);
            let mut auth = builder.init_auth();
            auth.set_origin(&self.local_domain);
        }

        let response = req.send().promise.await
            .map_err(|e| anyhow::anyhow!("federation proxy_resolve_user failed: {e}"))?;
        let key = response.get()
            .map_err(|e| anyhow::anyhow!("read proxy_resolve_user response: {e}"))?
            .get_identity_key()
            .map_err(|e| anyhow::anyhow!("get identity_key: {e}"))?;

        if key.is_empty() {
            Ok(None)
        } else {
            Ok(Some(key.to_vec()))
        }
    }

    /// Get an existing connection or create a new one to a peer domain.
    async fn get_or_connect(&self, domain: &str) -> anyhow::Result<quinn::Connection> {
        // Check for cached connection that's still alive.
        if let Some(conn) = self.connections.get(domain) {
            if conn.close_reason().is_none() {
                return Ok(conn.clone());
            }
        }

        let addr = self.peer_addresses.get(domain).ok_or_else(|| {
            anyhow::anyhow!("no federation peer configured for domain '{domain}'")
        })?;

        tracing::info!(domain = domain, addr = %addr, "connecting to federation peer");

        let conn = self
            .endpoint
            .connect(*addr, domain)
            .map_err(|e| anyhow::anyhow!("federation connect to {domain}: {e}"))?
            .await
            .with_context(|| format!("federation QUIC handshake with {domain}"))?;

        self.connections.insert(domain.to_string(), conn.clone());
        Ok(conn)
    }
}
