use anyhow::Result;
use async_trait::async_trait;
use futures_util::stream::{SplitSink, SplitStream, StreamExt};
use futures_util::SinkExt;
use serde_json::Value;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpStream;
use tokio::sync::mpsc::Sender;
use tokio::sync::Mutex;
use tokio::time::{timeout, Duration};
use tokio_tungstenite::tungstenite::Message as TokioMessage;
use tokio_tungstenite::WebSocketStream;

use crate::config::Config;
use crate::doslimit::{ConnectionLimits, GlobalLimits};
use crate::query::Query;
use crate::rpc::daemon::connection::Connection;
use crate::rpc::rpcstats::RpcStats;
use crate::signal::{NetworkNotifier, Waiter};
use crate::util::Channel;

use super::{communication::Communcation, server::ConnectionId, Message};
use crate::query::scripthash_subscriptions::ScripthashSubscriptions;

pub struct WsCommunication {
    writer: SplitSink<WebSocketStream<TcpStream>, TokioMessage>,
    reader: Option<SplitStream<WebSocketStream<TcpStream>>>,
}

impl WsCommunication {
    pub async fn init(stream: TcpStream) -> Result<Self> {
        let ws = tokio_tungstenite::accept_async(stream).await?;
        let (writer, reader) = ws.split();

        Ok(Self {
            reader: Some(reader),
            writer,
        })
    }
}

#[async_trait]
impl Communcation for WsCommunication {
    async fn send_values(&mut self, values: &[Value]) -> Result<()> {
        for value in values {
            let line = value.to_string();
            self.writer.feed(TokioMessage::Text(line.into())).await?;
        }
        self.writer.flush().await?;
        Ok(())
    }

    fn start_request_receiver(&mut self, rpc_queue_sender: Sender<Message>, idle_timeout: u64) {
        let reader = self.reader.take();

        crate::thread::spawn_task("ws_peer_reader", async move {
            // Give the reader to this task
            let mut reader = reader.unwrap();

            loop {
                Waiter::shutdown_check()?;
                let msg = timeout(Duration::from_secs(idle_timeout), reader.next()).await;

                let msg = match msg {
                    Ok(Some(message)) => message,
                    Ok(None) => break, // Connection closed
                    Err(_) => {
                        let _ = rpc_queue_sender.send(Message::Done).await;
                        return Err(anyhow!("Idle timeout {} exceeded", idle_timeout));
                    }
                };
                let msg = match msg {
                    Ok(msg) => msg,
                    Err(err) => {
                        trace!("ws peer error: {}", err);
                        let _ = rpc_queue_sender.send(Message::Done).await;
                        break;
                    }
                };

                if msg.is_text() {
                    let request = match msg.into_text() {
                        Ok(r) => r,
                        Err(e) => {
                            trace!("ws peer parse error: {}", e);
                            let _ = rpc_queue_sender.send(Message::Done).await;
                            break;
                        }
                    };
                    rpc_queue_sender
                        .send(Message::Request(request.to_string()))
                        .await?;
                } else if msg.is_binary() {
                    trace!("Ignoring binary message from ws peer");
                } else if msg.is_close() {
                    trace!("ws close event received");
                    let _ = rpc_queue_sender.send(Message::Done).await;
                    break;
                }
            }

            Ok(())
        });
    }

    async fn shutdown(&mut self) {
        let _ = self.writer.close().await;
    }
}

#[allow(clippy::too_many_arguments)]
pub async fn start_ws_server(
    conn_acceptor: Channel<Option<(TcpStream, SocketAddr)>>,
    global_limits: Arc<GlobalLimits>,
    config: Arc<Config>,
    query: Arc<Query>,
    stats: Arc<RpcStats>,
    connection_limits: ConnectionLimits,
    network_notifier: Arc<NetworkNotifier>,
    rpc_queue_senders: Arc<Mutex<std::collections::HashMap<ConnectionId, Sender<Message>>>>,
    subscription_manager: Arc<ScripthashSubscriptions>,
) -> Result<()> {
    while let Some(conn) = conn_acceptor.receiver().await.recv().await {
        let (mut stream, addr) = match conn {
            Some(c) => c,
            None => break,
        };

        let mut connections = match global_limits.inc_connection(&addr.ip()).await {
            Err(e) => {
                trace!("[{} ws] dropping peer - {}", addr, e);
                let _ = stream.shutdown().await;
                continue;
            }
            Ok(n) => n,
        };

        let config_cpy = config.clone();
        let query_cpy = query.clone();
        let stats_cpy = stats.clone();
        let limits_cpy = global_limits.clone();
        let network_notifier_cpy = network_notifier.clone();

        let senders_cpy = rpc_queue_senders.clone();
        let subscription_manager_cpy = subscription_manager.clone();
        crate::thread::spawn_task("ws_peer", async move {
            let rpc_queue: Channel<Message> = Channel::bounded(config_cpy.rpc_buffer_size);
            let conn_id = ConnectionId::new();

            let com = match WsCommunication::init(stream).await {
                Ok(c) => c,
                Err(e) => {
                    // Initialization failed, but we already incremented connection. Need to decrement.
                    if let Err(e) = limits_cpy.dec_connection(&addr.ip()).await {
                        warn!("Error decrementing connection count: {}", e);
                    }
                    return Err(e);
                }
            };

            senders_cpy.lock().await.insert(conn_id, rpc_queue.sender());
            debug!(
                "[{} ws] connected peer ({:?} out of {:?} connection slots used)",
                addr,
                connections,
                limits_cpy.connection_limits(),
            );

            let conn = Connection::new(
                config_cpy,
                query_cpy,
                addr,
                stats_cpy,
                connection_limits,
                Box::new(com),
                network_notifier_cpy,
                conn_id,
                subscription_manager_cpy.clone(),
            );
            let conn_id_for_cleanup = conn_id;
            conn.run(rpc_queue).await;
            subscription_manager_cpy
                .remove_connection(conn_id_for_cleanup)
                .await;
            senders_cpy.lock().await.remove(&conn_id_for_cleanup);

            match limits_cpy.dec_connection(&addr.ip()).await {
                Ok(n) => connections = n,
                Err(e) => error!("{}", e),
            };
            debug!(
                "[{} ws] disconnected peer ({:?} out of {:?} connection slots used)",
                addr,
                connections,
                limits_cpy.connection_limits(),
            );

            Ok(())
        });
    }

    info!("WebSocket connections no longer accepted");
    Ok(())
}
