vineyard

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

use std::net::SocketAddr;

use std::sync::Arc;

use std::time::Duration;

use anyhow::{Context as _, Result};

use futures_util::{Sink, SinkExt as _, StreamExt as _};

use nostr::event::Event;

use nostr::filter::{Filter, MatchEventOptions};

use nostr_database::{NostrDatabase as _, RejectedReason, SaveEventStatus};

use serde_json::{Value, json};

use tokio::io::{AsyncRead, AsyncWrite};

use tokio::sync::broadcast;

use tokio::time::timeout;

use tokio_tungstenite::WebSocketStream;

use tokio_tungstenite::tungstenite::{

    Error as WebSocketError, Message,

    protocol::{Role, WebSocketConfig},

};

use tokio_util::sync::CancellationToken;

use crate::config::Limits;

use crate::policy::{Policy, QueryPolicyResult, WritePolicyResult};

use crate::store::Store;

#[derive(Clone)]

pub struct Relay {

    store: Arc<Store>,

    policy: Policy,

    limits: Limits,

    events: broadcast::Sender<Event>,

    stop: CancellationToken,

}

struct Subscription {

    filters: Vec<Filter>,

    bytes: usize,

}

#[derive(Default)]

struct Session {

    subscriptions: HashMap<String, Subscription>,

}

impl Relay {

    pub(crate) fn new(

        store: Arc<Store>,

        policy: Policy,

        limits: Limits,

        stop: CancellationToken,

    ) -> Self {

        let (events, _) = broadcast::channel(limits.max_connections.max(32));

        Self {

            store,

            policy,

            limits,

            events,

            stop,

        }

    }

    pub(crate) async fn add_event(&self, event: Event) -> Result<SaveEventStatus> {

        event.verify().context("invalid event")?;

        let status = self.store.save_event(&event).await?;

        if status.is_success() {

            match self.events.send(event) {

                Ok(_) | Err(_) => {}

            }

        }

        Ok(status)

    }

    pub(crate) async fn take_connection<S>(&self, stream: S, address: SocketAddr) -> Result<()>

    where

        S: AsyncRead + AsyncWrite + Unpin,

    {

        let websocket = WebSocketStream::from_raw_socket(

            stream,

            Role::Server,

            Some(

                WebSocketConfig::default()

                    .max_message_size(Some(self.limits.max_message_bytes))

                    .max_frame_size(Some(self.limits.max_message_bytes)),

            ),

        )

        .await;

        let (mut sink, mut source) = websocket.split();

        let mut events = self.events.subscribe();

        let mut session = Session::default();

        loop {

            tokio::select! {

                biased;

                () = self.stop.cancelled() => return Ok(()),

                received = source.next() => {

                    let Some(message) = received else { return Ok(()); };

                    match message? {

                        Message::Text(text) => {

                            for response in self.handle_text(&mut session, address, &text).await {

                                self.send(&mut sink, response).await?;

                            }

                        }

                        Message::Ping(payload) => {

                            self.send_message(&mut sink, Message::Pong(payload)).await?;

                        }

                        Message::Close(_) => return Ok(()),

                        Message::Binary(_) | Message::Pong(_) | Message::Frame(_) => {}

                    }

                }

                received = events.recv() => {

                    let event = received.context("live event stream lagged")?;

                    let ids = session.subscriptions.iter()

                        .filter(|entry| entry.1.filters.iter().any(|filter| filter.match_event(&event, MatchEventOptions::new())))

                        .map(|entry| entry.0.clone())

                        .collect::<Vec<_>>();

                    for id in ids {

                        self.send(&mut sink, json!(["EVENT", id, event])).await?;

                    }

                }

            }

        }

    }

    async fn send<S>(&self, sink: &mut S, value: Value) -> Result<()>

    where

        S: Sink<Message, Error = WebSocketError> + Unpin,

    {

        self.send_message(sink, Message::Text(value.to_string().into()))

            .await

    }

    async fn send_message<S>(&self, sink: &mut S, message: Message) -> Result<()>

    where

        S: Sink<Message, Error = WebSocketError> + Unpin,

    {

        timeout(

            Duration::from_secs(self.limits.write_timeout_seconds),

            sink.send(message),

        )

        .await

        .context("websocket write timed out")??;

        Ok(())

    }

    async fn handle_text(

        &self,

        session: &mut Session,

        address: SocketAddr,

        text: &str,

    ) -> Vec<Value> {

        let Ok(value) = serde_json::from_str::<Value>(text) else {

            return vec![json!(["NOTICE", "invalid json"])];

        };

        let Some(frame) = value.as_array() else {

            return vec![json!(["NOTICE", "message must be an array"])];

        };

        match frame.first().and_then(Value::as_str) {

            Some("EVENT") => vec![self.handle_event(frame, address).await],

            Some("REQ") => self.handle_req(session, frame, address, text.len()).await,

            Some("CLOSE") => {

                if let Some(id) = frame.get(1).and_then(Value::as_str) {

                    session.subscriptions.remove(id);

                }

                Vec::new()

            }

            _ => vec![json!(["NOTICE", "unsupported message"])],

        }

    }

    async fn handle_event(&self, frame: &[Value], address: SocketAddr) -> Value {

        let Some(value) = frame.get(1) else {

            return json!(["NOTICE", "EVENT requires an event"]);

        };

        let id = value

            .get("id")

            .and_then(Value::as_str)

            .unwrap_or_default()

            .to_owned();

        let Ok(event) = serde_json::from_value::<Event>(value.clone()) else {

            return json!(["OK", id, false, "invalid: malformed event"]);

        };

        if event.verify().is_err() {

            return json!(["OK", event.id, false, "invalid: bad id or signature"]);

        }

        match self.policy.admit_event(&event, &address).await {

            WritePolicyResult::Accept => {}

            WritePolicyResult::Respond {

                prefix,

                message,

                status,

            } => return json!(["OK", event.id, status, format!("{prefix}: {message}")]),

        }

        let event_id = event.id;

        match self.add_event(event).await {

            Ok(SaveEventStatus::Success) => json!(["OK", event_id, true, ""]),

            Ok(SaveEventStatus::Rejected(RejectedReason::Duplicate)) => {

                json!(["OK", event_id, true, "duplicate: already stored"])

            }

            Ok(SaveEventStatus::Rejected(reason)) => {

                json!(["OK", event_id, false, format!("blocked: {reason:?}")])

            }

            Err(error) => {

                tracing::error!(event = %event_id, %error, "event storage failed");

                json!(["OK", event_id, false, "error: database unavailable"])

            }

        }

    }

    async fn handle_req(

        &self,

        session: &mut Session,

        frame: &[Value],

        address: SocketAddr,

        request_bytes: usize,

    ) -> Vec<Value> {

        let Some(id) = frame.get(1).and_then(Value::as_str) else {

            return vec![json!(["NOTICE", "REQ requires a subscription id"])];

        };

        if id.len() > 64 {

            return vec![closed(id, "blocked: subscription id is too long")];

        }

        if frame.len() < 3 {

            return vec![closed(id, "blocked: REQ requires a filter")];

        }

        if frame.len() - 2 > self.limits.max_filters {

            return vec![closed(id, "rate-limited: too many filters")];

        }

        if session.subscriptions.len() >= self.limits.max_subscriptions

            && !session.subscriptions.contains_key(id)

        {

            return vec![closed(id, "rate-limited: too many subscriptions")];

        }

        let old_bytes = session.subscriptions.get(id).map_or(0, |entry| entry.bytes);

        let active_bytes = session

            .subscriptions

            .values()

            .map(|entry| entry.bytes)

            .sum::<usize>()

            .saturating_sub(old_bytes);

        let max_subscription_bytes = self

            .limits

            .max_filter_bytes

            .saturating_mul(self.limits.max_subscriptions);

        if active_bytes.saturating_add(request_bytes) > max_subscription_bytes {

            return vec![closed(id, "rate-limited: subscription byte limit exceeded")];

        }

        let mut filters = Vec::with_capacity(frame.len() - 2);

        for value in &frame[2..] {

            let Ok(filter) = serde_json::from_value::<Filter>(value.clone()) else {

                return vec![closed(id, "blocked: invalid filter")];

            };

            match self.policy.admit_query(&filter, &address).await {

                QueryPolicyResult::Accept => filters.push(filter),

                QueryPolicyResult::Reject { prefix, message } => {

                    return vec![closed(id, &format!("{prefix}: {message}"))];

                }

            }

        }

        let mut events = BTreeSet::new();

        for filter in &filters {

            let implicit = filter.limit.is_none();

            let limit = filter.limit.unwrap_or(self.limits.max_query_results);

            let mut query = filter.clone();

            query.limit = Some(if implicit {

                limit.saturating_add(1)

            } else {

                limit

            });

            let found = match self.store.query(query).await {

                Ok(found) => found,

                Err(error) => {

                    let message = if error.to_string().contains("history byte limit exceeded") {

                        "rate-limited: history byte limit exceeded"

                    } else {

                        tracing::error!(%error, "history query failed");

                        "error: database unavailable"

                    };

                    return vec![closed(id, message)];

                }

            };

            if implicit && found.len() > limit {

                return vec![closed(id, "rate-limited: history result limit exceeded")];

            }

            events.extend(found);

            if events.len() > self.limits.max_query_results {

                return vec![closed(id, "rate-limited: too many results")];

            }

        }

        let bytes = events

            .iter()

            .map(|event| event.as_json().len())

            .sum::<usize>();

        if bytes > self.limits.max_query_bytes {

            return vec![closed(id, "rate-limited: history byte limit exceeded")];

        }

        let mut response = Vec::with_capacity(events.len() + 1);

        for event in events {

            if !event.is_expired() {

                response.push(json!(["EVENT", id, event]));

            }

        }

        response.push(json!(["EOSE", id]));

        session.subscriptions.insert(

            id.to_owned(),

            Subscription {

                filters,

                bytes: request_bytes,

            },

        );

        response

    }

}

fn closed(id: &str, message: &str) -> Value {

    json!(["CLOSED", id, message])

}