use std::{
    net::SocketAddr,
    sync::{atomic::AtomicBool, Arc},
};

use anyhow::Error;
use futures_util::{SinkExt, StreamExt};

use giterated_models::instance::Instance;

use giterated_models::object_backend::ObjectBackend;

use giterated_models::{
    authenticated::AuthenticatedPayload, message::GiteratedMessage, object::AnyObject,
    operation::AnyOperation,
};
use serde::Serialize;

use tokio::{net::TcpStream, sync::Mutex};
use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
use toml::Table;

use crate::{
    authentication::AuthenticationTokenGranter,
    backend::{MetadataBackend, RepositoryBackend, UserBackend},
    database_backend::DatabaseBackend,
    federation::connections::InstanceConnections,
    keys::PublicKeyCache,
};

use super::Connections;

pub async fn connection_wrapper(
    socket: WebSocketStream<TcpStream>,
    connections: Arc<Mutex<Connections>>,
    repository_backend: Arc<Mutex<dyn RepositoryBackend + Send>>,
    user_backend: Arc<Mutex<dyn UserBackend + Send>>,
    auth_granter: Arc<Mutex<AuthenticationTokenGranter>>,
    settings_backend: Arc<Mutex<dyn MetadataBackend + Send>>,
    addr: SocketAddr,
    instance: impl ToOwned<Owned = Instance>,
    instance_connections: Arc<Mutex<InstanceConnections>>,
    config: Table,
    backend: DatabaseBackend,
) {
    let connection_state = ConnectionState {
        socket: Arc::new(Mutex::new(socket)),
        connections,
        repository_backend,
        user_backend,
        auth_granter,
        settings_backend,
        addr,
        instance: instance.to_owned(),
        handshaked: Arc::new(AtomicBool::new(false)),
        key_cache: Arc::default(),
        instance_connections: instance_connections.clone(),
        config,
    };

    let _handshaked = false;

    loop {
        let mut socket = connection_state.socket.lock().await;
        let message = socket.next().await;
        drop(socket);

        match message {
            Some(Ok(message)) => {
                let payload = match message {
                    Message::Binary(payload) => payload,
                    Message::Ping(_) => {
                        let mut socket = connection_state.socket.lock().await;
                        let _ = socket.send(Message::Pong(vec![])).await;
                        drop(socket);
                        continue;
                    }
                    Message::Close(_) => return,
                    _ => continue,
                };

                let message: AuthenticatedPayload = bincode::deserialize(&payload).unwrap();

                let message: GiteratedMessage<AnyObject, AnyOperation> = message.into_message();

                backend
                    .object_operation(message.object, message.payload)
                    .await
                    .unwrap();
            }
            _ => {
                return;
            }
        }
    }

    // loop {
    //     let mut socket = connection_state.socket.lock().await;
    //     let message = socket.next().await;
    //     drop(socket);

    //     match message {
    //         Some(Ok(message)) => {
    //             let payload = match message {
    //                 Message::Binary(payload) => payload,
    //                 Message::Ping(_) => {
    //                     let mut socket = connection_state.socket.lock().await;
    //                     let _ = socket.send(Message::Pong(vec![])).await;
    //                     drop(socket);
    //                     continue;
    //                 }
    //                 Message::Close(_) => return,
    //                 _ => continue,
    //             };
    //             info!("one payload");

    //             let message = NetworkMessage(payload.clone());

    //             if !handshaked {
    //                 info!("im foo baring");
    //                 if handshake_handle(&message, &connection_state).await.is_ok() {
    //                     if connection_state.handshaked.load(Ordering::SeqCst) {
    //                         handshaked = true;
    //                     }
    //                 }
    //             } else {
    //                 let raw = serde_json::from_slice::<AuthenticatedPayload>(&payload).unwrap();

    //                 if let Some(target_instance) = &raw.target_instance {
    //                     if connection_state.instance != *target_instance {
    //                         // Forward request
    //                         info!("Forwarding message to {}", target_instance.url);
    //                         let mut instance_connections = instance_connections.lock().await;
    //                         let pool = instance_connections.get_or_open(&target_instance).unwrap();
    //                         let pool_clone = pool.clone();
    //                         drop(pool);

    //                         let result = wrap_forwarded(&pool_clone, raw).await;

    //                         let mut socket = connection_state.socket.lock().await;
    //                         let _ = socket.send(result).await;

    //                         continue;
    //                     }
    //                 }

    //                 let message_type = &raw.message_type;

    //                 info!("Handling message with type: {}", message_type);

    //                 match authentication_handle(message_type, &message, &connection_state).await {
    //                     Err(e) => {
    //                         let _ = connection_state
    //                             .send_raw(ConnectionError(e.to_string()))
    //                             .await;
    //                     }
    //                     Ok(true) => continue,
    //                     Ok(false) => {}
    //                 }

    //                 match repository_handle(message_type, &message, &connection_state).await {
    //                     Err(e) => {
    //                         let _ = connection_state
    //                             .send_raw(ConnectionError(e.to_string()))
    //                             .await;
    //                     }
    //                     Ok(true) => continue,
    //                     Ok(false) => {}
    //                 }

    //                 match user_handle(message_type, &message, &connection_state).await {
    //                     Err(e) => {
    //                         let _ = connection_state
    //                             .send_raw(ConnectionError(e.to_string()))
    //                             .await;
    //                     }
    //                     Ok(true) => continue,
    //                     Ok(false) => {}
    //                 }

    //                 match authentication_handle(message_type, &message, &connection_state).await {
    //                     Err(e) => {
    //                         let _ = connection_state
    //                             .send_raw(ConnectionError(e.to_string()))
    //                             .await;
    //                     }
    //                     Ok(true) => continue,
    //                     Ok(false) => {}
    //                 }

    //                 error!(
    //                     "Message completely unhandled: {}",
    //                     std::str::from_utf8(&payload).unwrap()
    //                 );
    //             }
    //         }
    //         Some(Err(e)) => {
    //             error!("Closing connection for {:?} for {}", e, addr);
    //             return;
    //         }
    //         _ => {
    //             info!("Unhandled");
    //             continue;
    //         }
    //     }
    // }
}

#[derive(Clone)]
pub struct ConnectionState {
    socket: Arc<Mutex<WebSocketStream<TcpStream>>>,
    pub connections: Arc<Mutex<Connections>>,
    pub repository_backend: Arc<Mutex<dyn RepositoryBackend + Send>>,
    pub user_backend: Arc<Mutex<dyn UserBackend + Send>>,
    pub auth_granter: Arc<Mutex<AuthenticationTokenGranter>>,
    pub settings_backend: Arc<Mutex<dyn MetadataBackend + Send>>,
    pub addr: SocketAddr,
    pub instance: Instance,
    pub handshaked: Arc<AtomicBool>,
    pub key_cache: Arc<Mutex<PublicKeyCache>>,
    pub instance_connections: Arc<Mutex<InstanceConnections>>,
    pub config: Table,
}

impl ConnectionState {
    pub async fn send<T: Serialize>(&self, message: T) -> Result<(), Error> {
        let payload = serde_json::to_string(&message)?;
        info!("Sending payload: {}", &payload);
        self.socket
            .lock()
            .await
            .send(Message::Binary(payload.into_bytes()))
            .await?;

        Ok(())
    }

    pub async fn send_raw<T: Serialize>(&self, message: T) -> Result<(), Error> {
        let payload = serde_json::to_string(&message)?;
        info!("Sending payload: {}", &payload);
        self.socket
            .lock()
            .await
            .send(Message::Binary(payload.into_bytes()))
            .await?;

        Ok(())
    }

    pub async fn public_key(&self, instance: &Instance) -> Result<String, Error> {
        let mut keys = self.key_cache.lock().await;
        keys.get(instance).await
    }
}