vineyard

use std::borrow::Cow;

use std::collections::{HashMap, HashSet};

use std::net::{IpAddr, SocketAddr};

use std::sync::{Arc, Mutex, PoisonError};

use std::time::{Duration, Instant};

use nostr::event::{Event, EventId, Kind};

use nostr::filter::Filter;

use nostr::key::PublicKey;

use nostr::message::MachineReadablePrefix;

use nostr::nips::nip65::extract_relay_list;

use nostr::types::Timestamp;

use nostr_database::{DatabaseEventStatus, NostrDatabase, error::Error};

use serde::Deserialize;

use crate::config::{Config, Limits};

#[derive(Debug)]

struct Usage {

    window: Instant,

    events: u32,

    queries: u32,

}

#[derive(Debug, Clone)]

pub struct Policy {

    database: Arc<dyn NostrDatabase>,

    limits: Limits,

    forge_keys: HashSet<PublicKey>,

    peers: Arc<Mutex<HashMap<IpAddr, Usage>>>,

}

#[derive(Deserialize)]

#[serde(deny_unknown_fields)]

struct Binding {

    version: u64,

    action: String,

    fingerprint: String,

    certificate: Option<String>,

    signature: Option<String>,

}

pub(crate) enum WritePolicyResult {

    Accept,

    Respond {

        prefix: MachineReadablePrefix,

        message: Cow<'static, str>,

        status: bool,

    },

}

impl WritePolicyResult {

    const fn accepted(prefix: MachineReadablePrefix, message: &'static str) -> Self {

        Self::Respond {

            prefix,

            message: Cow::Borrowed(message),

            status: true,

        }

    }

    const fn reject(prefix: MachineReadablePrefix, message: &'static str) -> Self {

        Self::Respond {

            prefix,

            message: Cow::Borrowed(message),

            status: false,

        }

    }

}

pub(crate) enum QueryPolicyResult {

    Accept,

    Reject {

        prefix: MachineReadablePrefix,

        message: Cow<'static, str>,

    },

}

impl QueryPolicyResult {

    const fn reject(prefix: MachineReadablePrefix, message: &'static str) -> Self {

        Self::Reject {

            prefix,

            message: Cow::Borrowed(message),

        }

    }

}

impl Policy {

    pub fn new(database: Arc<dyn NostrDatabase>, config: &Config) -> anyhow::Result<Self> {

        config.validate()?;

        Ok(Self {

            database,

            limits: config.limits.clone(),

            forge_keys: config.forge_keys()?,

            peers: Arc::new(Mutex::new(HashMap::new())),

        })

    }

    fn charge(&self, ip: IpAddr, publish: bool) -> bool {

        let mut peers = self.peers.lock().unwrap_or_else(PoisonError::into_inner);

        peers.retain(|_, usage| usage.window.elapsed() < Duration::from_secs(60));

        if peers.len() >= 4096 && !peers.contains_key(&ip) {

            return false;

        }

        let usage = peers.entry(ip).or_insert_with(|| Usage {

            window: Instant::now(),

            events: 0,

            queries: 0,

        });

        let (count, limit) = if publish {

            (&mut usage.events, self.limits.events_per_minute)

        } else {

            (&mut usage.queries, self.limits.queries_per_minute)

        };

        if *count >= limit {

            return false;

        }

        *count += 1;

        drop(peers);

        true

    }

    pub(crate) async fn admit_event(

        &self,

        event: &Event,

        address: &SocketAddr,

    ) -> WritePolicyResult {

        if !self.charge(address.ip(), true) {

            return WritePolicyResult::reject(

                MachineReadablePrefix::RateLimited,

                "publish rate exceeded",

            );

        }

        self.admit(event).await.unwrap_or_else(|error| {

            tracing::error!(%error, "event admission database failure");

            WritePolicyResult::reject(MachineReadablePrefix::Error, "database unavailable")

        })

    }

    pub(crate) async fn admit_peer(&self, event: &Event) -> WritePolicyResult {

        match self.admit(event).await {

            Ok(result) => result,

            Err(error) => {

                tracing::error!(%error, "event admission database failure");

                WritePolicyResult::reject(MachineReadablePrefix::Error, "database unavailable")

            }

        }

    }

    pub(crate) async fn admit_query(

        &self,

        filter: &Filter,

        address: &SocketAddr,

    ) -> QueryPolicyResult {

        if !self.charge(address.ip(), false) {

            return QueryPolicyResult::reject(

                MachineReadablePrefix::RateLimited,

                "query rate exceeded",

            );

        }

        if filter.as_json().len() > self.limits.max_filter_bytes {

            return QueryPolicyResult::reject(

                MachineReadablePrefix::RateLimited,

                "filter byte limit exceeded",

            );

        }

        if filter.search.is_some() {

            return QueryPolicyResult::reject(

                MachineReadablePrefix::Unsupported,

                "search is not supported",

            );

        }

        if let Some(limit) = filter.limit {

            return if limit > self.limits.max_query_results {

                QueryPolicyResult::reject(

                    MachineReadablePrefix::RateLimited,

                    "query result limit exceeded",

                )

            } else {

                QueryPolicyResult::Accept

            };

        }

        match self

            .database

            .count(filter.clone().limit(self.limits.max_query_results + 1))

            .await

        {

            Ok(count) if count <= self.limits.max_query_results => QueryPolicyResult::Accept,

            Ok(_) => QueryPolicyResult::reject(

                MachineReadablePrefix::RateLimited,

                "history result limit exceeded",

            ),

            Err(error) => {

                tracing::error!(%error, "query admission database failure");

                QueryPolicyResult::reject(MachineReadablePrefix::Error, "database unavailable")

            }

        }

    }

    async fn admit(&self, event: &Event) -> Result<WritePolicyResult, Error> {

        if let Some(result) = self.check_event(event) {

            return Ok(result);

        }

        match self.database.check_id(&event.id).await? {

            DatabaseEventStatus::Saved => {

                return Ok(WritePolicyResult::accepted(

                    MachineReadablePrefix::Duplicate,

                    "already stored",

                ));

            }

            DatabaseEventStatus::Deleted => return Ok(blocked("event is deleted")),

            DatabaseEventStatus::NotExistent => {}

        }

        self.admit_kind(event).await

    }

    fn check_event(&self, event: &Event) -> Option<WritePolicyResult> {

        if event.as_json().len() > self.limits.max_event_bytes

            || event.tags.len() > self.limits.max_tags

        {

            return Some(blocked("event exceeds size or tag limit"));

        }

        if event.created_at.as_secs() == 0

            || event.created_at.as_secs() > Timestamp::now().as_secs().saturating_add(900)

        {

            return Some(blocked("event timestamp is outside the supported range"));

        }

        if event.is_expired() || event.is_protected() {

            return Some(blocked("expired or protected events are not accepted"));

        }

        None

    }

    async fn admit_kind(&self, event: &Event) -> Result<WritePolicyResult, Error> {

        match event.kind.as_u16() {

            10002 => {

                if !self.forge_keys.contains(&event.pubkey)

                    && !self.registered_forge(event.pubkey).await?

                {

                    return Ok(blocked("relay list signer has no registered repository"));

                }

                let relay_tags = event

                    .tags

                    .iter()

                    .filter(|tag| tag.as_slice().first().is_some_and(|name| name == "r"))

                    .count();

                if relay_tags > 16 || extract_relay_list(event).count() != relay_tags {

                    return Ok(blocked("relay list is invalid or too large"));

                }

            }

            30617 => {

                let Some(identifier) = identifier(event) else {

                    return Ok(blocked("one valid repository identifier is required"));

                };

                if !self.registered(event.pubkey, identifier).await? {

                    let count = self

                        .database

                        .query(

                            Filter::new()

                                .kind(Kind::from(30617))

                                .limit(self.limits.max_repositories + 1),

                        )

                        .await?

                        .len();

                    if count >= self.limits.max_repositories {

                        return Ok(WritePolicyResult::reject(

                            MachineReadablePrefix::RateLimited,

                            "repository directory is full",

                        ));

                    }

                }

            }

            30618 => {

                let Some(identifier) = identifier(event) else {

                    return Ok(blocked("one valid repository identifier is required"));

                };

                if !self.registered(event.pubkey, identifier).await? {

                    return Ok(blocked("repository is not registered"));

                }

            }

            1617 | 1621 | 1630..=1632 => {

                let mut addresses = 0_usize;

                for tag in event

                    .tags

                    .iter()

                    .filter(|tag| tag.as_slice().first().is_some_and(|name| name == "a"))

                {

                    let Some((kind, author, identifier)) =

                        tag.as_slice().get(1).and_then(|value| coordinate(value))

                    else {

                        return Ok(blocked("invalid repository address"));

                    };

                    if kind != 30617 || !self.registered(author, identifier).await? {

                        return Ok(blocked("repository is not registered"));

                    }

                    if matches!(event.kind.as_u16(), 1630..=1632) && author != event.pubkey {

                        return Ok(blocked("status signer does not own the repository"));

                    }

                    addresses += 1;

                }

                if addresses != 1 {

                    return Ok(blocked("exactly one repository address is required"));

                }

            }

            30382 => {

                if !valid_binding(event) {

                    return Ok(blocked("invalid openpgp binding envelope"));

                }

            }

            5 => return self.deletion(event).await,

            _ => {

                return Ok(WritePolicyResult::reject(

                    MachineReadablePrefix::Unsupported,

                    "event kind is not supported",

                ));

            }

        }

        Ok(WritePolicyResult::Accept)

    }

    async fn registered(&self, author: PublicKey, identifier: &str) -> Result<bool, Error> {

        Ok(self

            .database

            .query(

                Filter::new()

                    .kind(Kind::from(30617))

                    .author(author)

                    .identifier(identifier)

                    .limit(1),

            )

            .await?

            .iter()

            .any(|event| !event.is_expired()))

    }

    async fn registered_forge(&self, author: PublicKey) -> Result<bool, Error> {

        Ok(self

            .database

            .query(

                Filter::new()

                    .kind(Kind::from(30617))

                    .author(author)

                    .limit(1),

            )

            .await?

            .iter()

            .any(|event| !event.is_expired()))

    }

    async fn deletion(&self, event: &Event) -> Result<WritePolicyResult, Error> {

        let mut targets = 0_usize;

        for tag in event.tags.iter() {

            let values = tag.as_slice();

            match values.first().map(String::as_str) {

                Some("e") => {

                    let Some(id) = values

                        .get(1)

                        .filter(|id| id.len() == 64)

                        .and_then(|id| id.parse::<EventId>().ok())

                    else {

                        return Ok(blocked("invalid deletion target"));

                    };

                    if let Some(target) = self.database.event_by_id(&id).await? {

                        if target.pubkey != event.pubkey || target.kind == Kind::EventDeletion {

                            return Ok(blocked("deletion target is not removable by this author"));

                        }

                    }

                    targets += 1;

                }

                Some("a") => {

                    let Some((kind, author, _)) = values.get(1).and_then(|value| coordinate(value))

                    else {

                        return Ok(blocked("invalid deletion address"));

                    };

                    if author != event.pubkey || !matches!(kind, 30617 | 30618 | 30382) {

                        return Ok(blocked("deletion address is not owned by this author"));

                    }

                    targets += 1;

                }

                _ => {}

            }

        }

        Ok(if targets == 0 {

            blocked("deletion target is required")

        } else {

            WritePolicyResult::Accept

        })

    }

}

const fn blocked(message: &'static str) -> WritePolicyResult {

    WritePolicyResult::reject(MachineReadablePrefix::Blocked, message)

}

fn valid_identifier(identifier: &str) -> bool {

    !identifier.is_empty() && identifier.len() <= 512 && !identifier.chars().any(char::is_control)

}

fn identifier(event: &Event) -> Option<&str> {

    let mut tags = event

        .tags

        .iter()

        .filter(|tag| tag.as_slice().first().is_some_and(|name| name == "d"));

    let tag = tags.next()?.as_slice();

    if tags.next().is_some() || tag.len() != 2 || !valid_identifier(&tag[1]) {

        return None;

    }

    Some(&tag[1])

}

pub(crate) fn coordinate(value: &str) -> Option<(u16, PublicKey, &str)> {

    let mut parts = value.splitn(3, ':');

    let kind = parts.next()?.parse().ok()?;

    let author_hex = parts.next()?;

    if author_hex.len() != 64 || !author_hex.bytes().all(|byte| byte.is_ascii_hexdigit()) {

        return None;

    }

    let author = author_hex.parse().ok()?;

    let identifier = parts.next()?;

    valid_identifier(identifier).then_some((kind, author, identifier))

}

fn valid_binding(event: &Event) -> bool {

    if event.content.len() > 192 * 1024 {

        return false;

    }

    let Ok(body) = serde_json::from_str::<Binding>(&event.content) else {

        return false;

    };

    if body.version != 1

        || !matches!(body.fingerprint.len(), 40 | 64)

        || !body

            .fingerprint

            .bytes()

            .all(|byte| byte.is_ascii_hexdigit())

        || identifier(event) != Some(body.fingerprint.as_str())

    {

        return false;

    }

    match body.action.as_str() {

        "bind" => {

            body.certificate

                .is_some_and(|certificate| !certificate.trim().is_empty())

                && body

                    .signature

                    .is_some_and(|signature| !signature.trim().is_empty())

        }

        "revoke" => body.certificate.is_none() && body.signature.is_none(),

        _ => false,

    }

}