vineyard

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

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use std::sync::Arc;

use std::time::{Duration, SystemTime, UNIX_EPOCH};

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

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

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

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

use nostr::key::PublicKey;

use nostr::message::MachineReadablePrefix;

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

use serde_json::{Value, json};

use tokio::net::{TcpStream, lookup_host};

use tokio::time::{Instant, sleep, timeout, timeout_at};

use tokio_tungstenite::tungstenite::{Message, protocol::WebSocketConfig};

use tokio_tungstenite::{

    MaybeTlsStream, WebSocketStream, client_async_tls_with_config, connect_async_with_config,

};

use tokio_util::sync::CancellationToken;

use tokio_util::task::TaskTracker;

use crate::config::{Config, Peer, relay_url};

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

use crate::relay::Relay;

use crate::store::Store;

const KINDS: [u16; 10] = [5, 1617, 1621, 1630, 1631, 1632, 10002, 30382, 30617, 30618];

const OVERLAP: u64 = 7 * 24 * 60 * 60;

const RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);

type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;

#[derive(Clone)]

struct Link {

    peer: Peer,

    config: Arc<Config>,

    store: Arc<Store>,

    relay: Relay,

    policy: Policy,

    discovered: bool,

    stop: CancellationToken,

}

struct Connection {

    socket: Socket,

    stop: CancellationToken,

    write_timeout: Duration,

    request_sequence: u64,

}

pub fn spawn(

    config: Arc<Config>,

    store: Arc<Store>,

    relay: Relay,

    policy: Policy,

    task_tracker: &TaskTracker,

    stop: CancellationToken,

) {

    let tasks = task_tracker.clone();

    task_tracker

        .clone()

        .spawn(manage(config, store, relay, policy, tasks, stop));

}

async fn manage(

    config: Arc<Config>,

    store: Arc<Store>,

    relay: Relay,

    policy: Policy,

    tasks: TaskTracker,

    stop: CancellationToken,

) {

    let mut active = HashMap::<String, CancellationToken>::new();

    loop {

        let mut desired = BTreeMap::new();

        for configured_peer in &config.peers {

            let mut normalized_peer = configured_peer.clone();

            let Some(url) = relay_url(&normalized_peer.url) else {

                continue;

            };

            normalized_peer.url = url.to_string();

            desired.insert(normalized_peer.url.clone(), (normalized_peer, false));

        }

        match store

            .journal

            .discovered(config.limits.max_discovered_peers)

            .await

        {

            Ok(urls) => {

                for raw in urls {

                    if let Some(peer) = discovered_peer(&config, &raw) {

                        desired.entry(peer.url.clone()).or_insert((peer, true));

                    }

                }

            }

            Err(error) => tracing::error!(%error, "loading discovered peers failed"),

        }

        active.retain(|url, peer_stop| {

            if desired.contains_key(url) {

                true

            } else {

                peer_stop.cancel();

                false

            }

        });

        for (url, (peer, discovered)) in desired {

            if active.contains_key(&url) {

                continue;

            }

            let peer_stop = stop.child_token();

            let link = Link {

                peer,

                config: Arc::clone(&config),

                store: Arc::clone(&store),

                relay: relay.clone(),

                policy: policy.clone(),

                discovered,

                stop: peer_stop.clone(),

            };

            tasks.spawn(run(link.clone(), true));

            tasks.spawn(run(link, false));

            active.insert(url, peer_stop);

        }

        tokio::select! {

            () = stop.cancelled() => break,

            () = sleep(Duration::from_secs(10)) => {}

        }

    }

    for peer_stop in active.into_values() {

        peer_stop.cancel();

    }

}

fn discovered_peer(config: &Config, raw: &str) -> Option<Peer> {

    let url = relay_url(raw)?;

    if (!config.allow_private_discovery && url.scheme() != "wss")

        || url == relay_url(&config.public_url)?

    {

        return None;

    }

    Some(Peer {

        url: url.to_string(),

        query_limit: config.limits.max_query_results,

        rescan_from: None,

    })

}

async fn run(link: Link, outbound: bool) {

    let direction = if outbound { "outbound" } else { "inbound" };

    let mut backoff = 5;

    while !link.stop.is_cancelled() {

        let started = Instant::now();

        let result = if outbound {

            publish(&link).await

        } else {

            subscribe(&link).await

        };

        if link.stop.is_cancelled() {

            break;

        }

        if let Err(error) = result {

            tracing::warn!(peer = %link.peer.url, %direction, %error, "peer connection interrupted");

        }

        if started.elapsed() > Duration::from_secs(60) {

            backoff = 5;

        }

        tokio::select! {

            () = link.stop.cancelled() => break,

            () = sleep(Duration::from_secs(backoff)) => {},

        }

        backoff = (backoff * 2).min(60);

    }

}

impl Connection {

    async fn open(link: &Link) -> Result<Self> {

        let config = WebSocketConfig::default()

            .max_message_size(Some(link.config.limits.max_message_bytes))

            .max_frame_size(Some(link.config.limits.max_message_bytes));

        let connected = async {

            if link.discovered {

                connect_discovered(link, config).await

            } else {

                connect_async_with_config(&link.peer.url, Some(config), false)

                    .await

                    .map(|(socket, _)| socket)

                    .map_err(Into::into)

            }

        };

        let socket = tokio::select! {

            () = link.stop.cancelled() => bail!("stopping"),

            result = timeout(RESPONSE_TIMEOUT, connected) => {

                result.context("peer connection timed out")??

            }

        };

        Ok(Self {

            socket,

            stop: link.stop.clone(),

            write_timeout: Duration::from_secs(link.config.limits.write_timeout_seconds),

            request_sequence: 0,

        })

    }

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

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

            .await

    }

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

        tokio::select! {

            () = self.stop.cancelled() => bail!("stopping"),

            result = timeout(self.write_timeout, self.socket.send(message)) => {

                result.context("peer write timed out")??;

                Ok(())

            }

        }

    }

    async fn next(&mut self, deadline: Instant) -> Result<Option<Value>> {

        loop {

            let message = tokio::select! {

                () = self.stop.cancelled() => bail!("stopping"),

                result = timeout_at(deadline, self.socket.next()) => match result {

                    Err(_) => return Ok(None),

                    Ok(Some(message)) => message?,

                    Ok(None) => bail!("peer disconnected"),

                }

            };

            match message {

                Message::Text(text) => return Ok(Some(serde_json::from_str(&text)?)),

                Message::Ping(payload) => self.send_message(Message::Pong(payload)).await?,

                Message::Close(_) => bail!("peer closed the connection"),

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

            }

        }

    }

}

async fn connect_discovered(link: &Link, config: WebSocketConfig) -> Result<Socket> {

    let url = url::Url::parse(&link.peer.url)?;

    let host = url.host_str().context("discovered peer has no host")?;

    let port = url

        .port_or_known_default()

        .context("discovered peer has no port")?;

    let mut last_error = None;

    for address in lookup_host((host, port)).await? {

        if !link.config.allow_private_discovery && !public_ip(address.ip()) {

            continue;

        }

        match TcpStream::connect(address).await {

            Ok(stream) => {

                return client_async_tls_with_config(&link.peer.url, stream, Some(config), None)

                    .await

                    .map(|(socket, _)| socket)

                    .map_err(Into::into);

            }

            Err(error) => last_error = Some(error),

        }

    }

    match last_error {

        Some(error) => Err(error.into()),

        None => bail!("discovered peer did not resolve to an allowed address"),

    }

}

fn public_ip(address: IpAddr) -> bool {

    match address {

        IpAddr::V4(ipv4) => public_ipv4(ipv4),

        IpAddr::V6(ipv6) => public_ipv6(ipv6),

    }

}

fn public_ipv4(ip: Ipv4Addr) -> bool {

    let [first, second, _, _] = ip.octets();

    !(ip.is_private()

        || ip.is_loopback()

        || ip.is_link_local()

        || ip.is_broadcast()

        || ip.is_documentation()

        || ip.is_unspecified()

        || ip.is_multicast()

        || first == 0

        || (first == 100 && (64..=127).contains(&second))

        || (first == 198 && matches!(second, 18 | 19)))

}

fn public_ipv6(ip: Ipv6Addr) -> bool {

    if let Some(ipv4) = ip.to_ipv4_mapped() {

        return public_ipv4(ipv4);

    }

    let segments = ip.segments();

    !(ip.is_loopback()

        || ip.is_unspecified()

        || ip.is_unique_local()

        || ip.is_unicast_link_local()

        || ip.is_multicast()

        || (segments[0] == 0x2001 && segments[1] == 0x0db8))

}

async fn publish(link: &Link) -> Result<()> {

    let mut connection = Connection::open(link).await?;

    let mut announced = HashSet::new();

    loop {

        let pending = link.store.journal.pending(&link.peer.url, 32).await?;

        if pending.is_empty() {

            connection

                .next(Instant::now() + Duration::from_secs(1))

                .await?;

            continue;

        }

        for (sequence, id) in pending {

            if announced.len() >= link.config.limits.max_repositories {

                announced.clear();

            }

            if let Some(event) = link.store.event_by_id(&id).await? {

                let repository_event = repository_coordinate(&event).is_some();

                if let Some(announcement) = announcement_for(link, &event).await? {

                    if announced.insert(announcement.id) {

                        if let Some(error) = publish_event(&mut connection, &announcement).await? {

                            announced.remove(&announcement.id);

                            defer(link, sequence, &error).await?;

                            continue;

                        }

                    }

                } else if repository_event {

                    link.store

                        .journal

                        .acknowledge(&link.peer.url, sequence)

                        .await?;

                    continue;

                }

                if let Some(error) = publish_event(&mut connection, &event).await? {

                    if link.store.event_by_id(&id).await?.is_some() {

                        defer(link, sequence, &error).await?;

                        continue;

                    }

                } else if event.kind == Kind::from(30617) {

                    announced.insert(event.id);

                }

            }

            link.store

                .journal

                .acknowledge(&link.peer.url, sequence)

                .await?;

        }

    }

}

async fn defer(link: &Link, sequence: u64, error: &str) -> Result<()> {

    let retry_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() + 300;

    link.store

        .journal

        .defer(&link.peer.url, sequence, error, retry_at)

        .await?;

    tracing::warn!(peer = %link.peer.url, sequence, %error, "peer delivery deferred");

    Ok(())

}

async fn publish_event(connection: &mut Connection, event: &Event) -> Result<Option<String>> {

    tokio::select! {

        () = connection.stop.cancelled() => bail!("stopping"),

        () = sleep(Duration::from_secs(1)) => {},

    }

    connection.send(json!(["EVENT", event])).await?;

    let id = event.id.to_hex();

    let deadline = Instant::now() + RESPONSE_TIMEOUT;

    loop {

        let frame = connection

            .next(deadline)

            .await?

            .context("peer acknowledgement timed out")?;

        if frame[0] == "OK" && frame[1].as_str() == Some(&id) {

            let reason = frame[3].as_str().unwrap_or("peer rejected event");

            return Ok(if frame[2] == true || reason.starts_with("duplicate:") {

                None

            } else {

                Some(reason.chars().take(512).collect())

            });

        }

    }

}

async fn announcement_for(link: &Link, event: &Event) -> Result<Option<Event>> {

    let Some((author, identifier)) = repository_coordinate(event) else {

        return Ok(None);

    };

    Ok(link

        .store

        .query(

            Filter::new()

                .kind(Kind::from(30617))

                .author(author)

                .identifier(identifier)

                .limit(1),

        )

        .await?

        .into_iter()

        .next())

}

fn repository_coordinate(event: &Event) -> Option<(PublicKey, String)> {

    if event.kind == Kind::from(30618) {

        event

            .tags

            .identifier()

            .map(|identifier| (event.pubkey, identifier))

    } else if matches!(event.kind.as_u16(), 1617 | 1621 | 1630..=1632) {

        event.tags.iter().find_map(|tag| {

            let values = tag.as_slice();

            if values.first().map(String::as_str) != Some("a") {

                return None;

            }

            let (kind, author, identifier) = coordinate(values.get(1)?)?;

            if kind != 30617 {

                return None;

            }

            Some((author, identifier.to_owned()))

        })

    } else {

        None

    }

}

async fn import(link: &Link, value: Value) -> Result<()> {

    let event: Event = serde_json::from_value(value)?;

    if event.verify().is_err() {

        tracing::debug!(peer = %link.peer.url, event = %event.id, "ignored invalid peer event");

        return Ok(());

    }

    match link.policy.admit_peer(&event).await {

        WritePolicyResult::Accept => {}

        WritePolicyResult::Respond { status: true, .. } => return Ok(()),

        WritePolicyResult::Respond {

            prefix: MachineReadablePrefix::Error | MachineReadablePrefix::RateLimited,

            message,

            ..

        } => bail!("peer import paused by {message}"),

        WritePolicyResult::Respond { message, .. } => {

            tracing::debug!(peer = %link.peer.url, event = %event.id, %message, "ignored out of scope peer event");

            return Ok(());

        }

    }

    match link.relay.add_event(event).await? {

        SaveEventStatus::Success

        | SaveEventStatus::Rejected(

            RejectedReason::Duplicate

            | RejectedReason::Deleted

            | RejectedReason::Expired

            | RejectedReason::Replaced

            | RejectedReason::InvalidDelete

            | RejectedReason::Vanished,

        ) => Ok(()),

        status @ SaveEventStatus::Rejected(_) => {

            bail!("peer event was not stored with status {status:?}")

        }

    }

}

async fn subscribe(link: &Link) -> Result<()> {

    let mut connection = Connection::open(link).await?;

    connection

        .send(json!(["REQ", "live", {"kinds": KINDS, "limit": 0_u8}]))

        .await?;

    let deadline = Instant::now() + RESPONSE_TIMEOUT;