package ch.usi.inf.sa4.sanmarinoes.smarthut.socket;

import ch.usi.inf.sa4.sanmarinoes.smarthut.config.GsonConfig;
import ch.usi.inf.sa4.sanmarinoes.smarthut.config.JWTTokenUtils;
import ch.usi.inf.sa4.sanmarinoes.smarthut.models.User;
import ch.usi.inf.sa4.sanmarinoes.smarthut.models.UserRepository;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.jsonwebtoken.ExpiredJwtException;
import java.io.IOException;
import java.util.Map;
import java.util.function.BiConsumer;
import javax.websocket.MessageHandler;
import javax.websocket.Session;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/** Generates MessageHandlers for unauthenticated socket sessions */
@Component
public class AuthenticationMessageListener {

    private Gson gson = GsonConfig.gson();

    private JWTTokenUtils jwtTokenUtils;

    private UserRepository userRepository;

    @Autowired
    public AuthenticationMessageListener(
            JWTTokenUtils jwtTokenUtils, UserRepository userRepository) {
        this.jwtTokenUtils = jwtTokenUtils;
        this.userRepository = userRepository;
    }

    /**
     * Generates a new message handler to handle socket authentication
     *
     * @param session the session to which authentication must be checked
     * @param authorizedSetter function to call once user is authenticated
     * @return a new message handler to handle socket authentication
     */
    MessageHandler.Whole<String> newHandler(
            final Session session, BiConsumer<User, Session> authorizedSetter) {
        return new MessageHandler.Whole<>() {
            @Override
            public void onMessage(final String message) {
                if (message == null) {
                    acknowledge(false);
                    return;
                }

                String token;
                String username;

                try {
                    token = gson.fromJson(message, JsonObject.class).get("token").getAsString();
                    username = jwtTokenUtils.getUsernameFromToken(token);
                } catch (ExpiredJwtException e) {
                    System.err.println(e.getMessage());
                    acknowledge(false);
                    return;
                } catch (Throwable ignored) {
                    System.out.println("Token format not valid");
                    acknowledge(false);
                    return;
                }

                final User user = userRepository.findByUsername(username);
                if (user == null || jwtTokenUtils.isTokenExpired(token)) {
                    System.out.println("Token not valid");
                    acknowledge(false);
                    return;
                }

                // Here user is authenticated
                session.removeMessageHandler(this);

                // Add user-session pair in authorized list
                authorizedSetter.accept(user, session);

                // update client to acknowledge authentication
                acknowledge(true);
            }

            private void acknowledge(boolean success) {
                try {
                    session.getBasicRemote()
                            .sendText(gson.toJson(Map.of("authenticated", success)));
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        };
    }
}