use crate::{subscription::ActiveSubscription, Error};
use futures_util::{SinkExt, StreamExt};
use nostr_rs_client::PoolSubscription;
use nostr_rs_types::{
    relayer::ROk,
    types::{Addr, SubscriptionId},
    Request, Response,
};
use std::{
    collections::HashMap,
    sync::atomic::{AtomicUsize, Ordering},
};
use tokio::{
    net::TcpStream,
    sync::{
        mpsc::{channel, Receiver, Sender},
        RwLock,
    },
    task::JoinHandle,
};
#[allow(unused_imports)]
use tokio_tungstenite::{accept_async, tungstenite::Message, WebSocketStream};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub struct ConnectionId(usize);

mod local;

pub use local::LocalConnection;

impl Default for ConnectionId {
    fn default() -> Self {
        static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
        Self(NEXT_ID.fetch_add(1, Ordering::SeqCst))
    }
}

impl ConnectionId {
    /// Connection ID for messages from the client pool or any special empty ConnectionId
    #[inline]
    pub fn new_empty() -> Self {
        Self(0)
    }

    /// Check if the connection id is empty
    ///
    /// Empty connection id is used for messages from Client pool to the relayer
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0 == 0
    }
}

type CompoundSubcription = (Option<PoolSubscription>, Vec<ActiveSubscription>);

#[derive(Debug)]
/// Relayer connection
///
/// The new connection struct. This struct spawn's a new worker that handles
/// upcoming messages form the client
pub struct Connection {
    conn_id: ConnectionId,
    sender: Sender<Response>,
    subscriptions: RwLock<HashMap<SubscriptionId, CompoundSubcription>>,
    handler: Option<JoinHandle<()>>,
}

const MAX_SUBSCRIPTIONS_BUFFER: usize = 100;

impl Drop for Connection {
    fn drop(&mut self) {
        if let Some(handler) = self.handler.take() {
            handler.abort();
        }
    }
}

impl Connection {
    /// Create a new local connection
    pub fn new_local_connection() -> (Self, Receiver<Response>) {
        let (sender, receiver) = channel(MAX_SUBSCRIPTIONS_BUFFER);
        (
            Self {
                conn_id: ConnectionId::default(),
                sender,
                subscriptions: Default::default(),
                handler: None,
            },
            receiver,
        )
    }

    /// Create new connection
    pub async fn new_connection(
        send_message_to_relayer: Sender<(ConnectionId, Request)>,
        disconnection_notify: Option<Sender<ConnectionId>>,
        stream: TcpStream,
    ) -> Result<Self, Error> {
        let websocket = accept_async(stream).await?;
        let conn_id = Default::default();
        let (sender, receiver) = channel(MAX_SUBSCRIPTIONS_BUFFER);
        Ok(Self {
            conn_id,
            sender,
            subscriptions: Default::default(),
            handler: Some(Self::spawn(
                send_message_to_relayer,
                websocket,
                receiver,
                disconnection_notify,
                conn_id,
            )),
        })
    }

    /// Spawn a new worker for this connection
    fn spawn(
        send_message_to_relayer: Sender<(ConnectionId, Request)>,
        websocket: WebSocketStream<TcpStream>,
        mut receiver: Receiver<Response>,
        disconnection_notify: Option<Sender<ConnectionId>>,
        conn_id: ConnectionId,
    ) -> JoinHandle<()> {
        tokio::spawn(async move {
            let (mut ws_sender, mut ws_receiver) = websocket.split();
            loop {
                tokio::select! {
                    Some(msg) = receiver.recv() => {
                        let msg = if let Ok(msg) = serde_json::to_string(&msg) {
                            msg
                        } else {
                            continue;
                        };
                        if let Err(err) =  ws_sender.send(Message::Text(msg)).await {
                            log::error!("Error sending message to client: {}", err);
                            break;
                        }
                    }
                    Some(msg) = ws_receiver.next() => {
                        if let Ok(Message::Text(msg)) = msg {
                            let msg: Result<Request, _> = serde_json::from_str(&msg);
                            match msg {
                                Ok(msg) => {
                                    if let Err(err) = send_message_to_relayer.send((conn_id, msg)).await {
                                        log::error!("Error sending message to relayer: {}", err);
                                    }
                                },
                                Err(err) => {
                                    log::error!("Error parsing message from client: {}", err);
                                    let reply: Response = ROk {
                                        id: Addr::default(),
                                        status: false,
                                        message: "Error parsing message".to_owned(),
                                    }.into();
                                    let reply = if let Ok(reply) = serde_json::to_string(&reply) {
                                        reply
                                    } else {
                                        continue;
                                    };
                                    if let Err(err) =  ws_sender.send(Message::Text(reply)).await {
                                        log::error!("Error sending message to client: {}", err);
                                        break;
                                    }
                                }
                            };
                        }
                    }
                    else => {
                        break;
                    }
                }
            }

            if let Some(disconnection_notify) = disconnection_notify {
                let _ = disconnection_notify.try_send(conn_id);
            }
        })
    }

    #[inline]
    /// Sends a message to this connection's websocket
    pub fn send(&self, response: Response) -> Result<(), Error> {
        self.sender
            .try_send(response)
            .map_err(|e| Error::TrySendError(Box::new(e)))
    }

    /// Get the sender for this connection
    pub fn get_sender(&self) -> Sender<Response> {
        self.sender.clone()
    }

    /// Get the connection id for this connection
    pub fn get_conn_id(&self) -> ConnectionId {
        self.conn_id
    }

    /// Create a subscription for this connection
    pub async fn subscribe(
        &self,
        id: SubscriptionId,
        subscriptions: (Option<PoolSubscription>, Vec<ActiveSubscription>),
    ) {
        self.subscriptions.write().await.insert(id, subscriptions);
    }

    /// Remove a subscription for this connection
    pub async fn unsubscribe(&self, id: &SubscriptionId) {
        self.subscriptions.write().await.remove(id);
    }
}
