/* This file is part of DarkFi (https://dark.fi)
 *
 * Copyright (C) 2020-2026 Dyne.org foundation
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::{collections::HashSet, future::Future, io::ErrorKind, sync::Arc};

use async_trait::async_trait;
use futures::future::try_join_all;
use parking_lot::Mutex as SyncMutex;
use smol::{
    io::{BufReader, ReadHalf, WriteHalf},
    lock::{Mutex, MutexGuard},
};
use tinyjson::JsonValue;
use tracing::{debug, info, warn};
use url::Url;

use super::{
    common::{
        http_read_from_stream_request, http_write_to_stream, read_from_stream, write_to_stream,
        INIT_BUF_SIZE,
    },
    jsonrpc::*,
    settings::RpcSettings,
};
use crate::{
    net::transport::{Listener, PtListener, PtStream},
    system::{StoppableTask, StoppableTaskPtr},
    util::logger::verbose,
    Error, Result,
};

/// Asynchronous trait implementing a handler for incoming JSON-RPC requests.
#[async_trait]
pub trait RequestHandler<T>: Sync + Send {
    async fn handle_request(&self, req: JsonRequest) -> JsonResult;

    async fn pong(&self, id: i64, _params: JsonValue) -> JsonResult {
        JsonResponse::new(JsonValue::String("pong".to_string()), id).into()
    }

    async fn connections_mut(&self) -> MutexGuard<'life0, HashSet<StoppableTaskPtr>>;

    async fn connections(&self) -> Vec<StoppableTaskPtr> {
        self.connections_mut().await.iter().cloned().collect()
    }

    async fn mark_connection(&self, task: StoppableTaskPtr) {
        self.connections_mut().await.insert(task);
    }

    async fn unmark_connection(&self, task: StoppableTaskPtr) {
        self.connections_mut().await.remove(&task);
    }

    async fn active_connections(&self) -> usize {
        self.connections_mut().await.len()
    }

    async fn stop_connections(&self) {
        info!(target: "rpc::server", "[RPC] Server stopped, closing connections");
        for (i, task) in self.connections().await.iter().enumerate() {
            debug!(target: "rpc::server", "Stopping connection #{i}");
            task.stop().await;
        }
    }
}

#[derive(Default)]
struct ConnectionTaskState {
    closing: bool,
    tasks: HashSet<StoppableTaskPtr>,
}

#[derive(Default)]
struct ConnectionTasks {
    state: SyncMutex<ConnectionTaskState>,
}

impl ConnectionTasks {
    /// Register and start a child while holding the task-set lock. This prevents
    /// a fast child from finishing before it has been registered and prevents
    /// new children from racing with connection shutdown.
    fn start<'a, MainFut>(
        self: &Arc<Self>,
        task: StoppableTaskPtr,
        main: MainFut,
        ex: Arc<smol::Executor<'a>>,
    ) where
        MainFut: Future<Output = Result<()>> + Send + 'a,
    {
        let mut state = self.state.lock();
        if state.closing {
            return
        }

        debug!(target: "rpc::server", "Adding background task {} to map", task.task_id);
        state.tasks.insert(task.clone());

        let tasks = self.clone();
        let task_ = task.clone();
        task.start(
            main,
            move |_| async move {
                debug!(
                    target: "rpc::server",
                    "Removing background task {} from map", task_.task_id,
                );
                tasks.state.lock().tasks.remove(&task_);
            },
            Error::DetachedTaskStopped,
            ex,
        );
    }

    fn close(&self) -> Vec<StoppableTaskPtr> {
        let mut state = self.state.lock();
        state.closing = true;
        state.tasks.iter().cloned().collect()
    }

    fn stop_all_nowait(&self) {
        for task in self.close() {
            task.stop_nowait();
        }
    }

    async fn stop_all(&self) {
        for task in self.close() {
            task.stop().await;
        }
        debug_assert!(self.state.lock().tasks.is_empty());
    }
}

struct ConnectionTasksGuard(Arc<ConnectionTasks>);

impl Drop for ConnectionTasksGuard {
    fn drop(&mut self) {
        self.0.stop_all_nowait();
    }
}

/// Auxiliary function to handle a request in the background.
async fn handle_request<T>(
    writer: Arc<Mutex<WriteHalf<Box<dyn PtStream>>>>,
    addr: Url,
    rh: Arc<impl RequestHandler<T> + 'static>,
    ex: Arc<smol::Executor<'_>>,
    tasks: Arc<ConnectionTasks>,
    settings: RpcSettings,
    req: JsonRequest,
) -> Result<()> {
    let req_id = req.id;

    // Handle disabled RPC methods
    let rep = if settings.is_method_disabled(&req.method) {
        debug!(target: "rpc::server", "RPC method {} is disabled", req.method);
        JsonError::new(ErrorCode::MethodNotFound, None, req.id).into()
    } else {
        rh.handle_request(req).await
    };

    match rep {
        JsonResult::Subscriber(subscriber) => {
            let task = StoppableTask::new();

            // Clone what needs to go in the background
            let addr_ = addr.clone();
            let writer_ = writer.clone();

            // Detach the subscriber so we can multiplex further requests
            tasks.start(
                task,
                async move {
                    // Subscribe to the inner method subscriber
                    let subscription = subscriber.publisher.subscribe().await;
                    loop {
                        // Listen for notifications
                        let notification = subscription.receive().await;

                        // Push notification
                        debug!(target: "rpc::server", "{addr_} <-- {}", notification.stringify().unwrap());
                        let notification = JsonResult::Notification(notification);

                        let mut writer_lock = writer_.lock().await;

                        #[allow(clippy::collapsible_else_if)]
                        if settings.use_http() {
                            if let Err(e) = http_write_to_stream(&mut writer_lock, &notification).await {
                                return Err(e.into())
                            }
                        } else {
                            if let Err(e) = write_to_stream(&mut writer_lock, &notification).await {
                                return Err(e.into())
                            }
                        }

                        drop(writer_lock);
                    }
                },
                ex.clone(),
            );
        }

        JsonResult::SubscriberWithReply(subscriber, reply) => {
            // Write the response
            debug!(target: "rpc::server", "{addr} <-- {}", reply.stringify()?);
            let mut writer_lock = writer.lock().await;
            if settings.use_http() {
                http_write_to_stream(&mut writer_lock, &reply.into()).await?;
            } else {
                write_to_stream(&mut writer_lock, &reply.into()).await?;
            }
            drop(writer_lock);

            let task = StoppableTask::new();
            // Clone what needs to go in the background
            let addr_ = addr.clone();
            let writer_ = writer.clone();

            // Detach the subscriber so we can multiplex further requests
            tasks.start(
                task,
                async move {
                    // Start the subscriber loop
                    let subscription = subscriber.publisher.subscribe().await;
                    loop {
                        // Listen for notifications
                        let notification = subscription.receive().await;

                        // Push notification
                        debug!(target: "rpc::server", "{addr_} <-- {}", notification.stringify().unwrap());
                        let notification = JsonResult::Notification(notification);

                        let mut writer_lock = writer_.lock().await;
                        #[allow(clippy::collapsible_else_if)]
                        if settings.use_http() {
                            if let Err(e) = http_write_to_stream(&mut writer_lock, &notification).await {
                                return Err(e.into())
                            }
                        } else {
                            if let Err(e) = write_to_stream(&mut writer_lock, &notification).await {
                                return Err(e.into())
                            }
                        }
                        drop(writer_lock);
                    }
                },
                ex.clone(),
            );
        }

        JsonResult::Request(_) | JsonResult::Notification(_) => {
            warn!(
                target: "rpc::server",
                "{addr}: handler returned Request/Notification for id={req_id}",
            );
            let err_rep: JsonResult = JsonError::new(
                ErrorCode::InternalError,
                Some("Handler returned a non-response variant".to_string()),
                req_id,
            )
            .into();

            let mut writer_lock = writer.lock().await;
            if settings.use_http() {
                http_write_to_stream(&mut writer_lock, &err_rep).await?;
            } else {
                write_to_stream(&mut writer_lock, &err_rep).await?;
            }
            drop(writer_lock);
        }

        JsonResult::Response(ref v) => {
            debug!(target: "rpc::server", "{addr} <-- {}", v.stringify()?);
            let mut writer_lock = writer.lock().await;
            if settings.use_http() {
                http_write_to_stream(&mut writer_lock, &rep).await?;
            } else {
                write_to_stream(&mut writer_lock, &rep).await?;
            }
            drop(writer_lock);
        }

        JsonResult::Error(ref v) => {
            debug!(target: "rpc::server", "{addr} <-- {}", v.stringify()?);
            let mut writer_lock = writer.lock().await;
            if settings.use_http() {
                http_write_to_stream(&mut writer_lock, &rep).await?;
            } else {
                write_to_stream(&mut writer_lock, &rep).await?;
            }
            drop(writer_lock);
        }
    }

    Ok(())
}

/// Accept function that should run inside a loop for accepting incoming
/// JSON-RPC requests and passing them to the [`RequestHandler`].
#[allow(clippy::type_complexity)]
async fn accept_with_tasks<'a, T: 'a>(
    reader: Arc<Mutex<BufReader<ReadHalf<Box<dyn PtStream>>>>>,
    writer: Arc<Mutex<WriteHalf<Box<dyn PtStream>>>>,
    addr: Url,
    rh: Arc<impl RequestHandler<T> + 'static>,
    tasks: Arc<ConnectionTasks>,
    settings: RpcSettings,
    ex: Arc<smol::Executor<'a>>,
) -> Result<()> {
    // Ensure cancellation signals all children even before the connection
    // task's stop handler gets a chance to await them.
    let _tasks_guard = ConnectionTasksGuard(tasks.clone());

    loop {
        let mut buf = Vec::with_capacity(INIT_BUF_SIZE);

        let mut reader_lock = reader.lock().await;
        if settings.use_http() {
            let _ = http_read_from_stream_request(&mut reader_lock, &mut buf).await?;
        } else {
            let _ = read_from_stream(&mut reader_lock, &mut buf).await?;
        }
        drop(reader_lock);

        let line = match String::from_utf8(buf) {
            Ok(v) => v,
            Err(e) => {
                warn!(
                    target: "rpc::server::accept",
                    "[RPC SERVER] Failed parsing string from read buffer: {e}"
                );
                return Err(e.into())
            }
        };

        // Parse the line as JSON
        let val: JsonValue = match line.trim().parse() {
            Ok(v) => v,
            Err(e) => {
                warn!(
                    target: "rpc::server::accept",
                    "[RPC SERVER] Failed parsing JSON string: {e}"
                );
                return Err(e.into())
            }
        };

        // Cast to JsonRequest
        let req = match JsonRequest::try_from(&val) {
            Ok(v) => v,
            Err(e) => {
                warn!(
                    target: "rpc::server::accept",
                    "[RPC SERVER] Failed casting JSON to a JsonRequest: {e}"
                );
                return Err(e.into())
            }
        };

        debug!(target: "rpc::server", "{addr} --> {}", val.stringify()?);

        // Create a new task to handle request in the background
        let task = StoppableTask::new();

        // Detach the task
        tasks.start(
            task,
            handle_request(
                writer.clone(),
                addr.clone(),
                rh.clone(),
                ex.clone(),
                tasks.clone(),
                settings.clone(),
                req,
            ),
            ex.clone(),
        );
    }
}

/// Accept incoming JSON-RPC requests and stop all request and subscriber tasks
/// before returning.
#[allow(clippy::type_complexity)]
pub async fn accept<'a, T: 'a>(
    reader: Arc<Mutex<BufReader<ReadHalf<Box<dyn PtStream>>>>>,
    writer: Arc<Mutex<WriteHalf<Box<dyn PtStream>>>>,
    addr: Url,
    rh: Arc<impl RequestHandler<T> + 'static>,
    settings: RpcSettings,
    ex: Arc<smol::Executor<'a>>,
) -> Result<()> {
    let tasks = Arc::new(ConnectionTasks::default());
    let result = accept_with_tasks(reader, writer, addr, rh, tasks.clone(), settings, ex).await;
    tasks.stop_all().await;
    result
}

/// Wrapper function around [`accept()`] to take the incoming connection and
/// pass it forward.
async fn run_accept_loop<'a, T: 'a>(
    listener: Box<dyn PtListener>,
    rh: Arc<impl RequestHandler<T> + 'static>,
    conn_limit: Option<usize>,
    settings: RpcSettings,
    ex: Arc<smol::Executor<'a>>,
) -> Result<()> {
    loop {
        let connection = match listener.next().await {
            Ok(negotiation) => negotiation.await,
            Err(err) => Err(err),
        };

        match connection {
            Ok((stream, url)) => {
                let rh_ = rh.clone();
                verbose!(target: "rpc::server", "[RPC] Server accepted conn from {url}");

                // Check and reserve a connection slot under the same lock so
                // simultaneous accepts on different listeners cannot exceed
                // the server-wide limit.
                let task = StoppableTask::new();
                let mut connections = rh.connections_mut().await;
                if let Some(limit) = conn_limit {
                    if connections.len() >= limit {
                        drop(connections);
                        debug!(
                            target: "rpc::server::run_accept_loop",
                            "[RPC] Connection limit ({limit}) reached, rejecting {url}",
                        );

                        // Send a JSONRPC error before dropping the stream so
                        // the client can tell "rejected" apart from "unreachabl;e".
                        let err: JsonResult = JsonError::new(
                            ErrorCode::ServerError(-32000),
                            Some("Server connection limit reached".to_string()),
                            0,
                        )
                        .into();

                        let (_, mut writer) = smol::io::split(stream);
                        if settings.use_http() {
                            let _ = http_write_to_stream(&mut writer, &err).await;
                        } else {
                            let _ = write_to_stream(&mut writer, &err).await;
                        }
                        // Writer drops here, closing the connection
                        continue
                    }
                }

                let (reader, writer) = smol::io::split(stream);
                let reader = Arc::new(Mutex::new(BufReader::new(reader)));
                let writer = Arc::new(Mutex::new(writer));

                let task_ = task.clone();
                let ex_ = ex.clone();
                let tasks = Arc::new(ConnectionTasks::default());
                let tasks_ = tasks.clone();

                // Register before starting so a connection that closes
                // immediately cannot finish before it is tracked.
                connections.insert(task.clone());
                task.clone().start(
                    accept_with_tasks(
                        reader,
                        writer,
                        url.clone(),
                        rh.clone(),
                        tasks,
                        settings.clone(),
                        ex_,
                    ),
                    |_| async move {
                        tasks_.stop_all().await;
                        verbose!(target: "rpc::server", "[RPC] Closed conn from {url}");
                        rh_.clone().unmark_connection(task_.clone()).await;
                    },
                    Error::ChannelStopped,
                    ex.clone(),
                );
                drop(connections);
            }

            // As per accept(2) recommendation:
            Err(e) if e.raw_os_error().is_some() => match e.raw_os_error().unwrap() {
                libc::EAGAIN | libc::ECONNABORTED | libc::EPROTO | libc::EINTR => continue,
                libc::ECONNRESET => {
                    warn!(
                        target: "rpc::server::run_accept_loop",
                        "[RPC] Connection reset by peer in accept_loop"
                    );
                    continue
                }
                libc::ETIMEDOUT => {
                    warn!(
                        target: "rpc::server::run_accept_loop",
                        "[RPC] Connection timed out in accept_loop"
                    );
                    continue
                }
                libc::EPIPE => {
                    warn!(
                        target: "rpc::server::run_accept_loop",
                        "[RPC] Broken pipe in accept_loop"
                    );
                    continue
                }
                x => {
                    warn!(
                        target: "rpc::server::run_accept_loop",
                        "[RPC] Unhandled OS Error: {e} {x}"
                    );
                    continue
                }
            },

            // In case a TLS handshake fails, we'll get this:
            Err(e) if e.kind() == ErrorKind::UnexpectedEof => continue,

            // Handle ErrorKind::Other
            Err(e) if e.kind() == ErrorKind::Other => {
                if let Some(inner) = std::error::Error::source(&e) {
                    if let Some(inner) = inner.downcast_ref::<futures_rustls::rustls::Error>() {
                        warn!(
                            target: "rpc::server::run_accept_loop",
                            "[RPC] rustls listener error: {inner:?}"
                        );
                        continue
                    }
                }

                warn!(
                    target: "rpc::server::run_accept_loop",
                    "[RPC] Unhandled ErrorKind::Other error: {e:?}"
                );
                continue
            }

            // Errors we didn't handle above:
            Err(e) => {
                warn!(
                    target: "rpc::server::run_accept_loop",
                    "[RPC] Unhandled listener.next() error: {e}"
                );
                continue
            }
        }
    }
}

/// Start a JSON-RPC server bound to the given accept URLs and use the
/// given [`RequestHandler`] to handle incoming requests.
///
/// The supported network schemes can be prefixed with `http+` to serve
/// JSON-RPC over HTTP/1.1.
pub async fn listen_and_serve<'a, T: 'a>(
    settings: RpcSettings,
    rh: Arc<impl RequestHandler<T> + 'static>,
    conn_limit: Option<usize>,
    ex: Arc<smol::Executor<'a>>,
) -> Result<()> {
    let mut listen_urls = Vec::with_capacity(settings.listen.len());
    let mut listener_settings = Vec::with_capacity(settings.listen.len());

    for endpoint in &settings.listen {
        let use_http = endpoint.scheme().starts_with("http+");
        let mut listen_url = endpoint.clone();
        if use_http {
            let scheme = endpoint.scheme().strip_prefix("http+").unwrap();
            listen_url.set_scheme(scheme).map_err(|_| Error::UrlParse(endpoint.to_string()))?;
        }

        listen_urls.push(listen_url);
        let mut endpoint_settings = settings.clone();
        endpoint_settings.listen = vec![endpoint.clone()];
        listener_settings.push(endpoint_settings);
    }

    // Bind every socket before starting any accept loop. If one address is
    // invalid or unavailable, no partially-running RPC server is left behind.
    let listeners = Listener::listen_all(listen_urls, None, false).await?;
    let accept_loops =
        listeners.into_iter().zip(listener_settings).map(|(listener, endpoint_settings)| {
            run_accept_loop(listener, rh.clone(), conn_limit, endpoint_settings, ex.clone())
        });

    try_join_all(accept_loops).await.map(|_| ())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        rpc::client::RpcClient,
        system::{msleep, Publisher},
    };
    use smol::{net::TcpListener, Executor};

    struct RpcServer {
        rpc_connections: Mutex<HashSet<StoppableTaskPtr>>,
        subscriber: JsonSubscriber,
    }

    #[async_trait]
    impl RequestHandler<()> for RpcServer {
        async fn handle_request(&self, req: JsonRequest) -> JsonResult {
            match req.method.as_str() {
                "ping" => return self.pong(req.id, req.params).await,
                "subscribe" => return self.subscriber.clone().into(),
                _ => panic!(),
            }
        }

        async fn connections_mut(&self) -> MutexGuard<'life0, HashSet<StoppableTaskPtr>> {
            self.rpc_connections.lock().await
        }
    }

    #[test]
    fn conn_manager() -> Result<()> {
        let executor = Arc::new(Executor::new());

        // This simulates a server and a client. Through the function, there
        // are some calls to sleep(), which are used for the tests, because
        // otherwise they execute too fast. In practice, The RPC server is
        // a long-running task so when polled, it should handle things in a
        // correct manner.
        smol::block_on(executor.run(async {
            // Find an available port
            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let sockaddr = listener.local_addr()?;
            let settings = RpcSettings {
                listen: vec![Url::parse(&format!("tcp://127.0.0.1:{}", sockaddr.port()))?],
                disabled_methods: vec![],
            };
            drop(listener);

            let rpc_server = Arc::new(RpcServer {
                rpc_connections: Mutex::new(HashSet::new()),
                subscriber: JsonSubscriber::new("event"),
            });
            let rpc_server_ = rpc_server.clone();

            let server_task = StoppableTask::new();
            server_task.clone().start(
                listen_and_serve(settings.clone(), rpc_server.clone(), None, executor.clone()),
                |res| async move {
                    match res {
                        Ok(()) | Err(Error::RpcServerStopped) => {
                            rpc_server_.stop_connections().await
                        }
                        Err(e) => panic!("{e}"),
                    }
                },
                Error::RpcServerStopped,
                executor.clone(),
            );

            // Let the server spawn
            msleep(500).await;

            // Connect a client
            let rpc_client0 = RpcClient::new(settings.listen[0].clone(), executor.clone()).await?;
            msleep(500).await;
            assert!(rpc_server.active_connections().await == 1);

            // Connect another client
            let rpc_client1 = RpcClient::new(settings.listen[0].clone(), executor.clone()).await?;
            msleep(500).await;
            assert!(rpc_server.active_connections().await == 2);

            // And another one
            let _rpc_client2 = RpcClient::new(settings.listen[0].clone(), executor.clone()).await?;
            msleep(500).await;
            assert!(rpc_server.active_connections().await == 3);

            // Close the first client
            rpc_client0.stop().await;
            msleep(500).await;
            assert!(rpc_server.active_connections().await == 2);

            // Close the second client
            rpc_client1.stop().await;
            msleep(500).await;
            assert!(rpc_server.active_connections().await == 1);

            // The Listener should be stopped when we stop the server task.
            server_task.stop().await;
            assert!(RpcClient::new(settings.listen[0].clone(), executor.clone()).await.is_err());

            // After the server is stopped, the connections tasks should also be stopped
            assert!(rpc_server.active_connections().await == 0);

            Ok(())
        }))
    }

    #[test]
    fn subscriber_tasks_follow_connection_lifetime() -> Result<()> {
        let executor = Arc::new(Executor::new());

        smol::block_on(executor.run(async {
            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let sockaddr = listener.local_addr()?;
            let settings = RpcSettings {
                listen: vec![Url::parse(&format!("tcp://127.0.0.1:{}", sockaddr.port()))?],
                disabled_methods: vec![],
            };
            drop(listener);

            let rpc_server = Arc::new(RpcServer {
                rpc_connections: Mutex::new(HashSet::new()),
                subscriber: JsonSubscriber::new("event"),
            });
            let rpc_server_ = rpc_server.clone();

            let server_task = StoppableTask::new();
            server_task.clone().start(
                listen_and_serve(settings.clone(), rpc_server.clone(), None, executor.clone()),
                |res| async move {
                    match res {
                        Ok(()) | Err(Error::RpcServerStopped) => {
                            rpc_server_.stop_connections().await
                        }
                        Err(e) => panic!("{e}"),
                    }
                },
                Error::RpcServerStopped,
                executor.clone(),
            );

            msleep(500).await;

            for _ in 0..32 {
                let client =
                    Arc::new(RpcClient::new(settings.listen[0].clone(), executor.clone()).await?);
                let client_ = client.clone();
                let subscriber_task = executor.spawn(async move {
                    client_
                        .subscribe(
                            JsonRequest::new("subscribe", JsonValue::Array(vec![])),
                            Publisher::new(),
                        )
                        .await
                });

                for _ in 0..100 {
                    if rpc_server.subscriber.publisher.active_subscriptions() == 1 {
                        break
                    }
                    msleep(10).await;
                }
                assert_eq!(rpc_server.subscriber.publisher.active_subscriptions(), 1);

                client.stop().await;
                assert!(subscriber_task.await.is_err());

                for _ in 0..100 {
                    if rpc_server.active_connections().await == 0 {
                        break
                    }
                    msleep(10).await;
                }
                assert_eq!(rpc_server.active_connections().await, 0);
                assert_eq!(rpc_server.subscriber.publisher.active_subscriptions(), 0);
            }

            server_task.stop().await;
            assert_eq!(rpc_server.active_connections().await, 0);
            assert_eq!(rpc_server.subscriber.publisher.active_subscriptions(), 0);

            Ok(())
        }))
    }
}