use crate::{
    logic::{ClientLogic, Error},
    login_state::LoginState,
};
use async_channel::Receiver;
use async_handle::Handle;
use async_rwlock::RwLock;
use basws_shared::{
    challenge, compression,
    protocol::{
        protocol_version, InstallationConfig, ServerRequest, ServerResponse, WsBatchResponse,
        WsRequest,
    },
    timing::current_timestamp,
    Version,
};
use futures::{
    stream::{SplitSink, SplitStream},
    SinkExt, StreamExt,
};
use once_cell::sync::OnceCell;
use std::{collections::HashMap, sync::Arc};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream};
use url::Url;

mod data;
mod waiter;
use self::{
    data::{ClientData, NetworkState},
    waiter::PendingResponse,
};

static REQUEST_COUNTER: OnceCell<Handle<u64>> = OnceCell::new();

pub struct Client<L>
where
    L: ClientLogic + ?Sized,
{
    data: Arc<ClientData<L>>,
}

impl<L> Clone for Client<L>
where
    L: ClientLogic,
{
    fn clone(&self) -> Self {
        Self {
            data: self.data.clone(),
        }
    }
}

impl<L> std::ops::Deref for Client<L>
where
    L: ClientLogic,
{
    type Target = L;

    fn deref(&self) -> &Self::Target {
        &self.data.logic
    }
}

impl<L> Client<L>
where
    L: ClientLogic + 'static,
{
    pub fn new(logic: L) -> Self {
        let (sender, receiver) = async_channel::unbounded();
        Self {
            data: Arc::new(ClientData {
                logic: Box::new(logic),
                sender,
                receiver,
                state: RwLock::new(NetworkState {
                    login_state: LoginState::Disconnected,
                    average_roundtrip: None,
                    average_server_timestamp_delta: None,
                }),
                mailboxes: RwLock::new(HashMap::new()),
            }),
        }
    }

    pub fn spawn(self) -> tokio::task::JoinHandle<anyhow::Result<()>> {
        tokio::spawn(self.run())
    }

    async fn set_login_state(&self, state: LoginState) -> anyhow::Result<()> {
        let client = self.clone();
        let mut network = self.data.state.write().await;
        network.login_state = state;
        self.data
            .logic
            .state_changed(&network.login_state, client)
            .await
    }

    pub async fn login_state(&self) -> LoginState {
        let network = self.data.state.read().await;
        network.login_state.clone()
    }

    async fn receiver(&self) -> Receiver<WsRequest<L::Request>> {
        self.data.receiver.clone()
    }

    pub async fn average_roundtrip(&self) -> Option<f64> {
        let network = self.data.state.read().await;
        network.average_roundtrip
    }

    pub async fn average_server_timestamp_delta(&self) -> Option<f64> {
        let network = self.data.state.read().await;
        network.average_server_timestamp_delta
    }

    pub async fn wait_for_response_to(
        &self,
        request_id: u64,
    ) -> Result<Vec<L::Response>, async_channel::RecvError> {
        let (sender, receiver) = async_channel::bounded(1);

        {
            let mut mailboxes = self.data.mailboxes.write().await;
            mailboxes.insert(request_id, sender);
        }

        receiver.recv().await
    }

    async fn server_url(&self) -> Url {
        self.data.logic.server_url()
    }

    async fn protocol_version(&self) -> Version {
        self.data.logic.protocol_version()
    }

    async fn stored_installation_config(&self) -> Option<InstallationConfig> {
        self.data.logic.stored_installation_config().await
    }

    async fn store_installation_config(&self, config: InstallationConfig) -> anyhow::Result<()> {
        self.data.logic.store_installation_config(config).await
    }

    async fn response_received(
        &self,
        response: L::Response,
        original_request_id: Option<u64>,
    ) -> anyhow::Result<()> {
        let client = self.clone();
        self.data
            .logic
            .response_received(response, original_request_id, client)
            .await
    }

    async fn handle_error(&self, error: Error) -> anyhow::Result<()> {
        let client = self.clone();
        self.data.logic.handle_error(error, client).await
    }

    pub async fn run(self) -> anyhow::Result<()> {
        loop {
            let url = self.server_url().await;

            let reconnect_delay = match connect_async(url).await {
                Ok((ws, _)) => {
                    let (tx, rx) = ws.split();
                    // TODO these errors should funnel through handle_error so that they can be gracefully ignored and the websocket run loop can continue
                    tokio::try_join!(self.send_loop(tx), self.receive_loop(rx))?;
                    None
                }
                Err(err) => {
                    self.handle_error(Error::Websocket(err)).await?;
                    Some(tokio::time::Duration::from_millis(500))
                }
            };

            self.set_login_state(LoginState::Disconnected).await?;

            if let Some(delay) = reconnect_delay {
                tokio::time::sleep(delay).await
            }
        }
    }

    async fn receive_loop<T>(&self, mut rx: SplitStream<WebSocketStream<T>>) -> anyhow::Result<()>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        loop {
            match rx.next().await {
                Some(Ok(Message::Binary(bytes))) => match compression::decompress(&bytes) {
                    Ok(response) => self.handle_batch_response(response).await?,
                    Err(error) => {
                        self.handle_error(Error::Compression(error)).await?;
                    }
                },
                Some(Err(err)) => return Err(anyhow::Error::from(err)),
                None => {
                    anyhow::bail!("Disconnected on read");
                }
                _ => {}
            }
        }
    }

    async fn handle_batch_response(
        &self,
        batch: WsBatchResponse<L::Response>,
    ) -> anyhow::Result<()> {
        if let Some(request_id) = batch.request_id {
            let mut mailboxes = self.data.mailboxes.write().await;
            if let Some(sender) = mailboxes.remove(&request_id) {
                let _ = sender
                    .send(
                        batch
                            .results
                            .iter()
                            .filter_map(|sr| {
                                if let ServerResponse::Response(response) = sr {
                                    Some(response)
                                } else {
                                    None
                                }
                            })
                            .cloned()
                            .collect(),
                    )
                    .await;
            }
        }

        for response in batch.results {
            match response {
                ServerResponse::NewInstallation(config) => {
                    self.set_login_state(LoginState::Connected {
                        installation_id: config.id,
                    })
                    .await?;
                    self.store_installation_config(config).await?;
                }
                ServerResponse::Connected { installation_id } => {
                    self.set_login_state(LoginState::Connected { installation_id })
                        .await?;
                }
                ServerResponse::Challenge { nonce } => {
                    let config = self.stored_installation_config().await.ok_or_else(|| {
                        anyhow::anyhow!("Server issued challenge, but client has no stored config")
                    })?;

                    self.server_request(ServerRequest::ChallengeResponse(
                        challenge::compute_challenge(&config.private_key, &nonce),
                    ))
                    .await?;
                }
                ServerResponse::Ping {
                    average_roundtrip,
                    average_server_timestamp_delta,
                    timestamp,
                } => {
                    let mut data = self.data.state.write().await;
                    data.average_roundtrip = average_roundtrip;
                    data.average_server_timestamp_delta = average_server_timestamp_delta;
                    self.server_request(ServerRequest::Pong {
                        original_timestamp: timestamp,
                        timestamp: current_timestamp(),
                    })
                    .await?;
                }
                ServerResponse::Error(error) => self.handle_error(Error::Server(error)).await?,
                ServerResponse::Response(response) => {
                    self.response_received(response, batch.request_id).await?
                }
            }
        }

        Ok(())
    }

    async fn server_request(&self, request: ServerRequest<L::Request>) -> anyhow::Result<u64> {
        let id = {
            let mut counter = REQUEST_COUNTER.get_or_init(|| Handle::new(0)).write().await;
            *counter = counter.wrapping_add(1);
            *counter
        };
        self.data.sender.send(WsRequest { id, request }).await?;
        Ok(id)
    }

    pub async fn request(&self, request: L::Request) -> anyhow::Result<PendingResponse<L>> {
        let request_id = self.server_request(ServerRequest::Request(request)).await?;
        Ok(PendingResponse::new(request_id, self.clone()))
    }

    async fn send_loop<T>(
        &self,
        mut tx: SplitSink<WebSocketStream<T>, Message>,
    ) -> anyhow::Result<()>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        let config = self.stored_installation_config().await;
        let installation_id = config.as_ref().map(|config| config.id);
        self.set_login_state(LoginState::Handshaking { config })
            .await?;

        self.server_request(ServerRequest::Greetings {
            protocol_version: protocol_version().to_string(),
            server_version: self.protocol_version().await.to_string(),
            installation_id,
        })
        .await?;

        let receiver = self.receiver().await;
        loop {
            let request = receiver.recv().await?;
            let bytes = compression::compress(&request);
            if let Err(err) = tx.send(Message::Binary(bytes)).await {
                return Err(anyhow::Error::from(err));
            }
        }
    }
}