/*
 * Copyright (C) 2025-2026  Tim Zimmermann <tim@aglais.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.aglais.zms;

import com.formdev.flatlaf.themes.FlatMacDarkLaf;
import com.formdev.flatlaf.themes.FlatMacLightLaf;
import com.formdev.flatlaf.util.SystemInfo;
import joptsimple.OptionParser;
import joptsimple.OptionSet;
import org.aglais.credentialstorage.SecretStore;
import org.aglais.credentialstorage.StorageProvider;
import org.aglais.credentialstorage.model.StoredCredential;
import org.aglais.zms.client.AglaisZMSClient;
import org.aglais.zms.client.ConnectionChangeReason;
import org.aglais.zms.client.ConnectionState;
import org.aglais.zms.client.IClientEventListener;
import org.aglais.zms.ui.LoginDialog;
import org.aglais.zms.ui.MainFrame;
import org.aglais.zms.ui.theme.OSThemeDetector;
import org.aglais.zms.util.ConfigUtil;
import org.aglais.zms.util.Log;
import org.aglais.zms.util.ResourceUtil;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.swing.*;
import java.awt.*;
import java.awt.event.WindowEvent;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.ResourceBundle;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.regex.Pattern;

import static org.aglais.zms.api.core.APIConstants.PROJECT_NAME_LOWERCASE_NO_SPACES;
import static org.aglais.zms.api.core.APIConstants.PROJECT_NAME_NO_SPACES;

/**
 * Main application class that handles initialization, database connectivity, theming and UI launch.
 */
public class AglaisZMS {

    private static final AglaisZMS INSTANCE = new AglaisZMS();
    private static final String LOG_TAG = AglaisZMS.class.getSimpleName();
    private static final String IV_PASSWORD_SEPARATOR = "$";

    private ClientConfig config;
    private MainFrame frame;
    private AglaisZMSClient client;
    private final ExecutorService backgroundExecutor = new ThreadPoolExecutor(1,
            Runtime.getRuntime().availableProcessors(), 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>());
    private List<BufferedImage> iconImages;
    private final SecretStore<StoredCredential> credentialStore = StorageProvider.getCredentialStorage(true,
            StorageProvider.SecureOption.REQUIRED);
    private SecretKey masterKey;

    /**
     * Application entry point
     *
     * @param args Command line arguments
     */
    public static void main(String[] args) {
        OptionParser parser = new OptionParser();
        String configDir;
        if (SystemInfo.isWindows)
            configDir = Paths.get(System.getenv("APPDATA"), PROJECT_NAME_NO_SPACES).toString();
        else
            configDir = Paths.get(System.getProperty("user.home"), ".config", PROJECT_NAME_LOWERCASE_NO_SPACES)
                    .toString();

        parser.accepts("config")
                .withRequiredArg()
                .ofType(String.class)
                .describedAs("Config file location")
                .defaultsTo(Paths.get(configDir, "client.json").toString());
        parser.accepts("help")
                .forHelp();

        OptionSet options = parser.parse(args);

        // Print only help if requested
        if (options.has("help")) {
            try {
                parser.printHelpOn(System.out);
            } catch (IOException e) {
                Log.e(LOG_TAG, "Failed to print help", e);
            }
            return;
        }

        INSTANCE.initialize((String) options.valueOf("config"));
        SwingUtilities.invokeLater(INSTANCE::start);
    }

    /**
     * Returns the singleton instance of the application
     *
     * @return The global AglaisZMS instance
     */
    public static AglaisZMS getInstance() {
        return INSTANCE;
    }

    /**
     * Returns the currently used {@link ClientConfig}
     *
     * @return Active client configuration
     */
    public ClientConfig getConfig() {
        return config;
    }

    /**
     * Returns the current database connection instance
     *
     * @return Active database connection, may be null if not connected
     */
    public AglaisZMSClient getClient() {
        return client;
    }

    /**
     * Returns the {@link ExecutorService} for running background tasks
     *
     * @return Background executor
     */
    public ExecutorService getBackgroundExecutor() {
        return backgroundExecutor;
    }

    /**
     * Attempts to establish a connection with the provided credentials
     *
     * @param dialog   The {@link LoginDialog} instance
     * @param url      API URL
     * @param username API user
     * @param password API password
     * @return true if connection successful, false otherwise
     */
    private boolean login(LoginDialog dialog, String url, String username, char[] password) {
        client = new AglaisZMSClient(url, username, password);

        Log.i(LOG_TAG, "API URL: " + url);
        Log.i(LOG_TAG, "Username: " + username);

        if (!client.connect()) {
            client = null;

            return false;
        }
        Log.i(LOG_TAG, "Successfully connected");

        client.registerListener((newState, reason) -> {
            if (newState == ConnectionState.DISCONNECTED && reason != ConnectionChangeReason.DISCONNECT_REQUEST) {
                SwingUtilities.invokeLater(() -> {
                    ResourceBundle uiBundle = ResourceBundle.getBundle("ui");

                    JOptionPane.showMessageDialog(dialog, uiBundle.getString("error_connection_lost"),
                            uiBundle.getString("error"), JOptionPane.ERROR_MESSAGE);
                    restart();
                });
            }
        });

        return true;
    }

    /**
     * Applies the system theme (light/dark) to the application UI
     *
     * @param isDark true to apply dark theme, false for light theme
     */
    private void applyTheme(boolean isDark) {
        try {
            UIManager.setLookAndFeel(isDark ? new FlatMacDarkLaf() : new FlatMacLightLaf());
            Log.d(LOG_TAG, "Successfully applied " + (isDark ? "Dark" : "Light") + " theme");
        } catch (UnsupportedLookAndFeelException e) {
            Log.e(LOG_TAG, "Failed to apply " + (isDark ? "Dark" : "Light") + " theme", e);
        }

        for (Window window : Window.getWindows()) {
            SwingUtilities.updateComponentTreeUI(window);
        }
    }

    /**
     * Retrieves a list of app icon images for multiple resolutions.
     * The resolutions included are 16, 32, 64, 96, and 128.
     *
     * @return a list of {@link BufferedImage}s corresponding to each resolution; entries may be null if images are missing
     */
    public List<BufferedImage> getIconImages() {
        if (iconImages != null)
            return iconImages;

        iconImages = new ArrayList<>();

        for (int resolution : List.of(16, 32, 64, 96, 128)) {
            iconImages.add(ResourceUtil.getImage(ResourceUtil.APP_ICON_NAME + "_" + resolution + "."
                    + ResourceUtil.APP_ICON_EXTENSION));
        }

        return iconImages;
    }

    /**
     * Get the master key for encryption/decryption of passwords.
     * If there is no key stored in the credential store yet, a new one will be generated and stored.
     *
     * @return Master {@link SecretKey}
     */
    private SecretKey getOrCreateMasterKey() {
        if (credentialStore == null) {
            Log.w(LOG_TAG, "Secure credential store not available");

            return null;
        }

        if (masterKey != null)
            return masterKey;

        ResourceBundle uiBundle = ResourceBundle.getBundle("ui");
        String credentialName = uiBundle.getString("program_name");
        String keyName = "MasterKey";

        StoredCredential credential = credentialStore.get(credentialName);
        if (credential == null) {
            Log.i(LOG_TAG, "Generating new master key");

            try {
                KeyGenerator keyGen = KeyGenerator.getInstance("AES");

                keyGen.init(256);
                masterKey = keyGen.generateKey();
            } catch (NoSuchAlgorithmException e) {
                Log.e(LOG_TAG, "Failed to generate master key", e);

                return null;
            }

            credential = new StoredCredential(keyName, Base64.getEncoder().encodeToString(masterKey.getEncoded())
                    .toCharArray());
            if (!credentialStore.add(credentialName, credential)) {
                Log.e(LOG_TAG, "Failed to store master key");

                masterKey = null;
            }
            credential.clear();
        } else {
            byte[] keyBytes = Base64.getDecoder().decode(new String(credential.getPassword()));
            masterKey = new SecretKeySpec(keyBytes, "AES");

            credential.clear();
        }

        return masterKey;
    }

    /**
     * Initialize {@link Cipher} for encrypting/decrypting passwords
     *
     * @param key    The {@link SecretKey} to use
     * @param iv     The IV to use
     * @param opmode The operation mode, can be {@link Cipher#ENCRYPT_MODE} or {@link Cipher#DECRYPT_MODE}
     * @return Successfully initialized {@link Cipher}, null if an error occurred
     */
    private Cipher initPasswordCipher(SecretKey key, byte[] iv, int opmode) {
        try {
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec spec = new GCMParameterSpec(128, iv);

            cipher.init(opmode, key, spec);

            return cipher;
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidAlgorithmParameterException |
                 InvalidKeyException e) {
            Log.e(LOG_TAG, "Failed to initialize cipher", e);
        }

        return null;
    }

    /**
     * Encrypt a password for securely storing it.
     * The result is <b>_base64 encoded IV used for encryption_$_AES encrypted password_</b>
     *
     * @param password The password to encrypt
     * @return Encrypted password string, null if unsupported or an error occurred
     */
    public String encryptPassword(char[] password) {
        SecretKey key = getOrCreateMasterKey();

        if (key == null || password == null)
            return null;

        byte[] iv = new byte[12];
        try {
            SecureRandom.getInstanceStrong().nextBytes(iv);
        } catch (NoSuchAlgorithmException e) {
            Log.e(LOG_TAG, "Failed to generate IV", e);

            return null;
        }

        Cipher cipher = initPasswordCipher(key, iv, Cipher.ENCRYPT_MODE);
        if (cipher == null)
            return null;

        byte[] cipherText;
        try {
            cipherText = cipher.doFinal(new String(password).getBytes(StandardCharsets.UTF_8));
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            Log.e(LOG_TAG, "Failed to encrypt password", e);

            return null;
        }

        return Base64.getEncoder().encodeToString(iv) + IV_PASSWORD_SEPARATOR
                + Base64.getEncoder().encodeToString(cipherText);
    }

    /**
     * Decrypt an encrypted password.
     * The format is expected to match the one used by {@link #encryptPassword}.
     *
     * @param encrypted The encrypted password to decrypt
     * @return Decrypted password string, null if unsupported or an error occurred
     */
    public char[] decryptPassword(String encrypted) {
        SecretKey key = getOrCreateMasterKey();

        if (key == null || encrypted == null)
            return null;

        String[] parts = encrypted.split(Pattern.quote(IV_PASSWORD_SEPARATOR));
        if (parts.length != 2) {
            Log.e(LOG_TAG, "Invalid encrypted password (len=" + parts.length + ")");

            return null;
        }
        String ivString = parts[0];
        String cipherTextString = parts[1];
        byte[] iv = Base64.getDecoder().decode(ivString);
        byte[] cipherText = Base64.getDecoder().decode(cipherTextString);

        Cipher cipher = initPasswordCipher(key, iv, Cipher.DECRYPT_MODE);
        if (cipher == null)
            return null;

        byte[] plainTextBytes;
        try {
            plainTextBytes = cipher.doFinal(cipherText);
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            Log.e(LOG_TAG, "Failed to decrypt password", e);

            return null;
        }

        return new String(plainTextBytes, StandardCharsets.UTF_8).toCharArray();
    }

    /**
     * Initializes the application
     * - Loads the client configuration
     * - Sets up system theme detection and application
     *
     * @param configPath The path to load the config file from
     */
    private void initialize(String configPath) {
        config = ConfigUtil.loadConfig(configPath, ClientConfig.class, ClientConfig::new);
        config.configPath = configPath;

        OSThemeDetector themeDetector = OSThemeDetector.getInstance();

        applyTheme(themeDetector.isDark());
        themeDetector.registerListener((isDark) -> SwingUtilities.invokeLater(() -> applyTheme(isDark)));

        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            Log.i(LOG_TAG, "Shutting down");
            stop();
        }));
    }

    /**
     * Starts the application
     * - Shows login dialog
     * - Handles database version checking and upgrades
     * - Launches the main application frame if login is successful
     */
    public void start() {
        Log.i(LOG_TAG, "Starting login dialog");
        LoginDialog loginDialog = new LoginDialog(this, this::login);
        ResourceBundle uiBundle = ResourceBundle.getBundle("ui");

        loginDialog.setIconImages(getIconImages());
        loginDialog.setTitle(uiBundle.getString("dialog_login_title"));
        loginDialog.setAlwaysOnTop(true);
        loginDialog.setResizable(false);
        loginDialog.start();

        if (client != null) {
            Log.i(LOG_TAG, "Checking server version");
            int serverVersion = client.getVersion();
            boolean serverSupported = serverVersion >= client.getMinSupportedServerVersion();
            boolean clientSupported = client.getCurrentClientVersion() >= client.getMinSupportedClientVersion();
            if (!serverSupported || !clientSupported) {
                JOptionPane.showMessageDialog(loginDialog, uiBundle.getString(!clientSupported
                                ? "error_outdated_client" : "error_outdated_server"), uiBundle.getString("error"),
                        JOptionPane.ERROR_MESSAGE);

                System.exit(0);
            }

            Log.i(LOG_TAG, "Starting main frame");
            frame = new MainFrame(this);
            frame.start();
        } else {
            Log.i(LOG_TAG, "Not logged in");
            System.exit(0);
        }
    }

    /**
     * Deinitializes the application
     * - Terminates the main window
     * - Disconnects the client
     */
    public void stop() {
        if (client != null)
            client.disconnect();
    }

    /**
     * Restarts the application by closing the window, calling {@link #stop()} followed by {@link #start()}
     */
    public void restart() {
        if (frame != null) {
            frame.setExitOnClose(false);
            frame.getFrame().dispatchEvent(new WindowEvent(frame.getFrame(), WindowEvent.WINDOW_CLOSING));
        }
        stop();
        start();
    }
}
