vineyard

use std::collections::HashMap;

use std::convert::Infallible;

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

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

use std::time::Duration;

use anyhow::Context as _;

use http_body_util::Full;

use hyper::body::{Bytes, Incoming};

use hyper::header::{ACCEPT, ACCESS_CONTROL_ALLOW_ORIGIN, CONTENT_TYPE, HeaderValue, UPGRADE};

use hyper::server::conn::http1::Builder as Http1Builder;

use hyper::service::service_fn;

use hyper::upgrade::on as upgrade;

use hyper::{Method, Request, Response, StatusCode};

use hyper_util::rt::{TokioIo, TokioTimer};

use nostr_database::NostrDatabase;

use serde_json::json;

use tokio::net::TcpListener;

use tokio::sync::{OwnedSemaphorePermit, Semaphore};

use tokio::task::{JoinHandle, spawn as spawn_task};

use tokio::time::timeout;

use tokio_tungstenite::tungstenite::handshake::server::create_response;

use tokio_util::sync::CancellationToken;

use tokio_util::task::TaskTracker;

use crate::config::Config;

use crate::federation::spawn as spawn_federation;

use crate::policy::Policy;

use crate::relay::Relay;

use crate::store::Store;

type Body = Full<Bytes>;

struct ConnectionLease {

    ip: IpAddr,

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

    _permit: OwnedSemaphorePermit,

}

impl Drop for ConnectionLease {

    fn drop(&mut self) {

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

        if let Some(count) = peers.get_mut(&self.ip) {

            *count -= 1;

            if *count == 0 {

                peers.remove(&self.ip);

            }

        }

    }

}

#[derive(Clone)]

struct HttpState {

    relay: Relay,

    store: Arc<Store>,

    config: Arc<Config>,

    websocket_path: String,

    tasks: TaskTracker,

    stop: CancellationToken,

}

pub struct RunningRelay {

    address: SocketAddr,

    path: String,

    state: HttpState,

    listener: JoinHandle<anyhow::Result<()>>,

}

impl RunningRelay {

    #[must_use]

    pub fn url(&self) -> String {

        format!("ws://{}{}", self.address, self.path)

    }

    pub async fn shutdown(self) -> anyhow::Result<()> {

        self.state.stop.cancel();

        self.listener.await.context("joining listener")??;

        self.state.tasks.close();

        timeout(Duration::from_secs(30), self.state.tasks.wait())

            .await

            .context("relay connections did not drain")?;

        Ok(())

    }

}

pub async fn spawn(config: Config) -> anyhow::Result<RunningRelay> {

    config.validate()?;

    let store = Store::open(&config).await?;

    let policy_database: Arc<dyn NostrDatabase> = Arc::<Store>::clone(&store);

    let policy = Policy::new(policy_database, &config)?;

    let listener = TcpListener::bind(config.listen)

        .await

        .context("binding relay listener")?;

    let address = listener.local_addr()?;

    let path = config.websocket_path()?;

    let stop = CancellationToken::new();

    let relay = Relay::new(

        Arc::clone(&store),

        policy.clone(),

        config.limits.clone(),

        stop.clone(),

    );

    let state = HttpState {

        relay,

        store,

        config: Arc::new(config),

        websocket_path: path.clone(),

        tasks: TaskTracker::new(),

        stop,

    };

    let task_state = state.clone();

    spawn_federation(

        Arc::clone(&state.config),

        Arc::clone(&state.store),

        state.relay.clone(),

        policy,

        &state.tasks,

        state.stop.clone(),

    );

    let listener_task = spawn_task(async move { accept_connections(listener, task_state).await });

    tracing::info!(%address, %path, "relay is listening");

    Ok(RunningRelay {

        address,

        path,

        state,

        listener: listener_task,

    })

}

async fn accept_connections(listener: TcpListener, state: HttpState) -> anyhow::Result<()> {

    let slots = Arc::new(Semaphore::new(state.config.limits.max_connections));

    let peers = Arc::new(Mutex::new(HashMap::<IpAddr, usize>::new()));

    loop {

        let (stream, address) = tokio::select! {

            biased;

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

            accepted = listener.accept() => accepted?,

        };

        let Ok(permit) = Arc::clone(&slots).try_acquire_owned() else {

            continue;

        };

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

        let count = peer_counts.entry(address.ip()).or_default();

        if *count >= state.config.limits.max_connections_per_ip {

            continue;

        }

        *count += 1;

        drop(peer_counts);

        let lease = Arc::new(ConnectionLease {

            ip: address.ip(),

            peers: Arc::clone(&peers),

            _permit: permit,

        });

        let handler_state = state.clone();

        let stop = state.stop.clone();

        state.tasks.spawn(async move {

            let handler = service_fn(move |request| {

                handle_request(

                    request,

                    address,

                    handler_state.clone(),

                    Arc::clone(&lease),

                )

            });

            let mut builder = Http1Builder::new();

            builder.timer(TokioTimer::new()).header_read_timeout(Duration::from_secs(10));

            let connection = builder.serve_connection(TokioIo::new(stream), handler).with_upgrades();

            tokio::pin!(connection);

            tokio::select! {

                result = &mut connection => {

                    if let Err(error) = result { tracing::debug!(%address, %error, "http connection ended"); }

                }

                () = stop.cancelled() => {}

            }

        });

    }

}

async fn handle_request(

    mut request: Request<Incoming>,

    address: SocketAddr,

    state: HttpState,

    lease: Arc<ConnectionLease>,

) -> Result<Response<Body>, Infallible> {

    if state.stop.is_cancelled() {

        return Ok(response(

            StatusCode::SERVICE_UNAVAILABLE,

            "text/plain",

            "relay is stopping",

        ));

    }

    if request.method() != Method::GET {

        return Ok(response(

            StatusCode::METHOD_NOT_ALLOWED,

            "text/plain",

            "GET required",

        ));

    }

    if request.uri().path() == "/healthz" {

        let healthy = timeout(Duration::from_secs(2), state.store.healthy())

            .await

            .unwrap_or(false);

        return Ok(response(

            if healthy {

                StatusCode::OK

            } else {

                StatusCode::SERVICE_UNAVAILABLE

            },

            "text/plain",

            if healthy {

                "ok"

            } else {

                "database unavailable"

            },

        ));

    }

    if request.uri().path() != state.websocket_path {

        return Ok(response(StatusCode::NOT_FOUND, "text/plain", "not found"));

    }

    if !request.headers().contains_key(UPGRADE) {

        let wants_information = request

            .headers()

            .get(ACCEPT)

            .and_then(|value| value.to_str().ok())

            .is_some_and(|value| {

                value.split(',').any(|media_type| {

                    media_type.trim().split(';').next() == Some("application/nostr+json")

                })

            });

        return Ok(if wants_information {

            response(

                StatusCode::OK,

                "application/nostr+json",

                &information(&state.config),

            )

        } else {

            response(

                StatusCode::OK,

                "text/plain",

                "vineyard\na persistent nostr relay for git forges\n",

            )

        });

    }

    let upgrade = upgrade(&mut request);

    let Ok(handshake) = create_response(&request.map(|_body| ())) else {

        return Ok(response(

            StatusCode::BAD_REQUEST,

            "text/plain",

            "invalid WebSocket handshake",

        ));

    };

    let relay = state.relay.clone();

    state.tasks.spawn(async move {

        let _lease = lease;

        if let Ok(Ok(stream)) = timeout(Duration::from_secs(10), upgrade).await {

            if let Err(error) = relay.take_connection(TokioIo::new(stream), address).await {

                tracing::debug!(%address, %error, "websocket connection ended");

            }

        } else {

            tracing::debug!(%address, "websocket upgrade did not complete");

        }

    });

    Ok(handshake.map(|()| Full::new(Bytes::new())))

}

fn response(status: StatusCode, content_type: &str, body: &str) -> Response<Body> {

    let mut response = Response::new(Full::new(Bytes::copy_from_slice(body.as_bytes())));

    *response.status_mut() = status;

    if let Ok(value) = HeaderValue::from_str(content_type) {

        response.headers_mut().insert(CONTENT_TYPE, value);

    }

    response

        .headers_mut()

        .insert(ACCESS_CONTROL_ALLOW_ORIGIN, HeaderValue::from_static("*"));

    response

}

fn information(config: &Config) -> String {

    json!({

        "name": "vineyard",

        "description": "a persistent nostr relay for git forges",

        "software": "vineyard",

        "version": env!("CARGO_PKG_VERSION"),

        "supported_nips": [1, 9, 11, 34, 65],

        "forge_pubkeys": config.forge_pubkeys,

        "limitation": {

            "max_message_length": config.limits.max_message_bytes,

            "max_event_tags": config.limits.max_tags,

            "max_subscriptions": config.limits.max_subscriptions,

            "max_filters": config.limits.max_filters,

            "max_limit": config.limits.max_query_results,

            "max_subid_length": 64,

            "auth_required": false,

            "payment_required": false,

            "restricted_writes": true

        }

    })

    .to_string()

}