use std::collections::HashSet;
use std::sync::{Arc, RwLock, RwLockReadGuard};

use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use utoipa::ToSchema;

use crate::transport::EncodedMessage;
use crate::transport::{MessageHandler, ParsedMessage, Transport};

#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
pub struct FirewallRules {
    /// Messages from these peers are dropped before the application sees them.
    #[serde(default)]
    pub inbound: HashSet<String>,
    /// Messages to these peers are silently dropped, like a lossy network.
    #[serde(default)]
    pub outbound: HashSet<String>,
}

/// Uses a std lock: checks are tiny and the guard is never held across an await.
#[derive(Default)]
pub struct Firewall {
    rules: RwLock<FirewallRules>,
}

impl Firewall {
    fn read(&self) -> RwLockReadGuard<'_, FirewallRules> {
        self.rules
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    pub fn rules(&self) -> FirewallRules {
        self.read().clone()
    }

    pub fn set_rules(&self, rules: FirewallRules) {
        *self
            .rules
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = rules;
    }

    pub fn blocks_inbound(&self, peer: &str) -> bool {
        self.read().inbound.contains(peer)
    }

    pub fn blocks_outbound(&self, peer: &str) -> bool {
        self.read().outbound.contains(peer)
    }
}

pub struct FirewallTransport {
    inner: Arc<dyn Transport>,
    firewall: Arc<Firewall>,
}

impl FirewallTransport {
    pub fn new(inner: Arc<dyn Transport>, firewall: Arc<Firewall>) -> Self {
        Self { inner, firewall }
    }
}

#[async_trait]
impl Transport for FirewallTransport {
    fn local_id(&self) -> &str {
        self.inner.local_id()
    }

    async fn resolve_peer(&self, hint: &str) -> Result<String> {
        self.inner.resolve_peer(hint).await
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        self.inner.encode(peer, message)
    }

    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        if self.firewall.blocks_outbound(&message.peer) {
            tracing::trace!(peer = %message.peer, "Firewall dropped outbound message");
            return Ok(());
        }
        self.inner.send_encoded(message).await
    }

    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        let firewall = Arc::clone(&self.firewall);

        self.inner
            .start_consumer(Box::new(move |sender, message| {
                if firewall.blocks_inbound(&sender) {
                    tracing::trace!(peer = %sender, "Firewall dropped inbound message");
                    return Box::pin(async { Ok(()) });
                }
                handler(sender, message)
            }))
            .await
    }

    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        self.inner.wait_for_exit().await
    }
}