package app;

import jakarta.jms.Destination;
import jakarta.jms.JMSConsumer;
import jakarta.jms.JMSProducer;
import jakarta.jms.JMSContext;
import jakarta.jms.JMSException;
import jakarta.jms.TextMessage;

import com.ibm.msg.client.jakarta.jms.JmsConnectionFactory;
import com.ibm.msg.client.jakarta.jms.JmsFactoryFactory;
import com.ibm.msg.client.jakarta.wmq.WMQConstants;
import com.ibm.mq.constants.MQConstants;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpRequest.BodyPublishers;
import java.net.http.HttpResponse;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;

public class mq {

    // MQ Queue Manager Configuration.
    private static String queueManager = "";
    private static String queueName = "";
    private static String channel = "";
    private static int port = 0;

    // Cognito JWKS Configuration.
    private static String username = "";
    private static String password = "";
    private static String clientId = "";
    private static String region = "";

    /*
        Purpose: This Java application demonstrates how to send and receive messages to/from an IBM MQ queue using
        AWS Cognito for JWKS-based authentication.

        Dependencies:
        - Jakarta Messaging API (JMS) for MQ interaction.
        - IBM MQ Java client libraries.
        - AWS SDK for Java (specifically, HttpClient for Cognito token retrieval).

        Configuration:
        Environment variables for MQ connection details (queue manager, queue name, channel, port) and Cognito configuration details
        (username, password, client ID, region) are loaded from the `env.json` file.

        Flow:
        - The application loads environment variables from env.json.
        - Retrives an access token from AWS Cognito using the provided credentials and the HTTP package.
        - Connects to the MQ queue and sends/receives messages specified by the arguments.
        - After successful message exchange, it prints a success message and exits.
        */

    public static void main(String[] args) {
        System.out.println("\nJava Jakarta App to send message to MQ with AWS Cognito providing JWKS Authentication.\n");

        // Load environment variables from env.json.
        getEnvVars();

        // Obtain access token from Cognito.
        String accessToken = getCognitoAccessToken();

        System.out.println("\nConnecting to MQ...");
        if (args.length > 0) {
            // Connect to MQ and use the command line arguments to send custom messages.
            connectAndTestMQ(accessToken, args);
        } else {
            // Use default message if no arguments are provided
            System.out.println("No arg messages found, using default test message.");
            String[] msgs = {"Hello from MQ using AWS Cognito JWKS Token Authtentication!"};
            // Connect to MQ using the defualt message.
            connectAndTestMQ(accessToken, msgs);
        }
        System.out.println("\nSuccessfully completed MQ Cognito JWKS validation test.");
        // Expected end of execution.
        System.exit(0);
    }

    // connectAndTestMQ, establishes a connection to an MQ queue using provided credentials and sends/receives messages.
    private static void connectAndTestMQ(String accessToken, String[] msgs){
        JMSContext context = null;
        Destination destination = null;
        JMSProducer producer = null;
        JMSConsumer consumer = null;

        try {
            System.out.println("Creating connection factory.");
            JmsFactoryFactory ff = JmsFactoryFactory.getInstance(WMQConstants.JAKARTA_WMQ_PROVIDER);
            JmsConnectionFactory cf = ff.createConnectionFactory();

            // Set MQ connection properties.
            cf.setIntProperty(WMQConstants.WMQ_CONNECTION_MODE, WMQConstants.WMQ_CM_CLIENT);
            cf.setIntProperty(WMQConstants.WMQ_PORT, port);
            cf.setStringProperty(WMQConstants.WMQ_CHANNEL, channel);
            cf.setStringProperty(WMQConstants.WMQ_QUEUE_MANAGER, queueManager);
            // Use the Cognito access token for authentication.
            cf.setStringProperty(WMQConstants.PASSWORD, accessToken);

            context = cf.createContext();
            destination = context.createQueue("queue:///" + queueName);

            System.out.println("Successfully connected to MQ Queue Manager and created queue.");
            System.out.println("\nPutting and Getting messages to the Queue...");

            for (String msg: msgs){
                producer = context.createProducer();
                consumer = context.createConsumer(destination);

                // Send a message
                TextMessage sendMessage = context.createTextMessage(msg);
                producer.send(destination, sendMessage);
                System.out.println("Sent message: " + sendMessage.getText());

                // Recieve a message
                TextMessage receivedMessage = (TextMessage) consumer.receive(10000);
                if (receivedMessage != null){
                    System.out.println("Received message: " + receivedMessage.getText());
                } else {
                    System.out.println("error: no message received?!");
                    System.exit(1);
                }
            }
        } catch (JMSException ex) {
            System.out.println("error JMSException occurred: " + ex.toString());
            System.exit(1);
        }
    }

    // getCognitoAccessToken, retrieves an access token from AWS Cognito using provided credentials.
    // It constructs a request, sends it to Cognito, and extracts the access token from the response.
    private static String getCognitoAccessToken(){
        System.out.println("\nGetting user access token from Amazon AWS Cognito...");

        String accessToken = "";
        String requestData = "{" +
            "  \"AuthParameters\": {" +
            "    \"USERNAME\": \"" + username + "\"," +
            "    \"PASSWORD\": \"" + password + "\"" +
            "  }," +
            "  \"AuthFlow\": \"USER_PASSWORD_AUTH\"," +
            "  \"ClientId\": \"" + clientId + "\"" +
            "}";
        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
            .uri(URI.create("https://cognito-idp." + region + ".amazonaws.com/"))
            .POST(BodyPublishers.ofString(requestData))
            .setHeader("X-Amz-Target", "AWSCognitoIdentityProviderService.InitiateAuth")
            .setHeader("Content-Type", "application/x-amz-json-1.1")
            .build();

        System.out.println("Cognito Request: " + request + requestData);

        try{
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            String JWT = response.body();

            Pattern accessTokenPattern = Pattern.compile("\"AccessToken\"\s*:\s*\"([^,]*)\",");
            Matcher accessTokenMatcher = accessTokenPattern.matcher(JWT);

            if (accessTokenMatcher.find()) {
                accessToken = accessTokenMatcher.group(1);
                System.out.println("Cognito Access Token: " + accessToken);
            } else {
                System.out.println("error: Could not parse JWT token response");
                System.exit(1);
            }
        } catch(IOException ex){
            System.out.println(ex.toString());
            System.exit(1);
        } catch (InterruptedException e) {
            System.out.println(e.toString());
            System.exit(1);
        }
        return accessToken;
    }

    // getEnvVars, loads environment variables from env.json.
    private static void getEnvVars(){
        String filename = "env.json";
        System.out.printf("Loading environment variables from %s...\n", filename);
        String envs = readFile(filename);

        System.out.println("MQ Config:");
        queueManager = getVar(envs, "QMGR");
        queueName = getVar(envs, "QUEUE_NAME");
        channel = getVar(envs, "CHANNEL");

        String portStr = getVar(envs, "PORT");
        try{
            port = Integer.parseInt(portStr);
        } catch (NumberFormatException e) {
            System.out.printf("error: port value provided [%s] is not a number: %s", portStr, e.toString());
            System.exit(1);
        }

        System.out.println("\nCognito Config:");
        username = getVar(envs, "USERNAME");
        password = getVar(envs, "PASSWORD");
        clientId = getVar(envs, "CLIENT_ID");
        region = getVar(envs, "REGION");
    }

    // getVar, retrieves the value of a specific environment variable from a JSON-formatted string.
    // If the variable is found, it returns its value; otherwise, it prints an error message and exits the program.
    private static String getVar(String envs, String name){
        Pattern pattern = Pattern.compile("\"" + name + "\"\s*:\s*\"([^,]*)\"");
        Matcher matcher = pattern.matcher(envs);

        if (matcher.find()) {
            String var = matcher.group(1);

            if (var == ""){
                System.out.println("error: value of " + name + ", is a empty string.");
                System.exit(1);
            }

            System.out.printf(" - %s: \t%s\n", name, var);
            return var;
        } else {
            System.out.println("error: Could not find env var " + name + ", check JSON format is correct and variable isn't missing.");
            System.exit(1);
        }
        return "";
    }

    // Readfile, reads contents of filename and retuns the file content as a String.
    // If error occurs, prints the eror and exits the program.
    private static String readFile(String filename) {
		try {
			return new String(Files.readAllBytes(Paths.get(filename)));
		} catch (IOException e) {
            System.out.printf("error: failed to read file [%s]: %s", filename, e.toString());
            System.exit(1);
            return null;
		}
	}
}

echo "Creating MQ Jakarta Application and installing dependencies..."

# Tidy existing files.
rm -rf MQClient

# Make the necessary directories and install the relevant java dependencies.
mkdir MQClient
cd MQClient
curl -o com.ibm.mq.jakarta.client-9.4.2.0.jar https://repo1.maven.org/maven2/com/ibm/mq/com.ibm.mq.jakarta.client/9.4.2.0/com.ibm.mq.jakarta.client-9.4.2.0.jar
curl -o jakarta.jms-api-3.1.0.jar https://repo1.maven.org/maven2/jakarta/jms/jakarta.jms-api/3.1.0/jakarta.jms-api-3.1.0.jar
curl -o json-20250107.jar https://repo1.maven.org/maven2/org/json/json/20250107/json-20250107.jar
mkdir app

# Create the json env file.
cat <<EOF > env.json
{
    "MQ_CONFIG": {
        "QMGR": "QM1",
        "QUEUE_NAME": "DEV.QUEUE.1",
        "CHANNEL": "DEV.APP.SVRCONN",
        "PORT": "1414"
    },
    "COGNITO_CONFIG": {
        "USERNAME": "app",
        "PASSWORD": "P@ssw0rd1",
        "CLIENT_ID": "",
        "REGION": "eu-west-2"
    }
}
EOF

echo "env.json created with the following content:"
cat env.json

# Create the Makefile.
cat <<EOF > Makefile
.PHONY: run
run: build run-app

.PHONY: build
build:
	javac -cp ./com.ibm.mq.jakarta.client-9.4.2.0.jar:./jakarta.jms-api-3.1.0.jar:./json-20250107.jar app/mq.java

.PHONY: run-app
run-app:
	java -cp ./com.ibm.mq.jakarta.client-9.4.2.0.jar:./jakarta.jms-api-3.1.0.jar:./json-20250107.jar:. app.mq \$(m1) \$(m2) \$(m3) \$(m4) \$(m5)

.PHONY: clean
clean:
	rm -f app/mq.class
EOF

echo "Makefile created with the following content:"
cat Makefile

# Create the Java Jakarta MQ app.
cd app
cat <<EOF > mq.java
package app;

import jakarta.jms.Destination;
import jakarta.jms.JMSConsumer;
import jakarta.jms.JMSProducer;
import jakarta.jms.JMSContext;
import jakarta.jms.JMSException;
import jakarta.jms.TextMessage;

import com.ibm.msg.client.jakarta.jms.JmsConnectionFactory;
import com.ibm.msg.client.jakarta.jms.JmsFactoryFactory;
import com.ibm.msg.client.jakarta.wmq.WMQConstants;
import com.ibm.mq.constants.MQConstants;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpRequest.BodyPublishers;
import java.net.http.HttpResponse;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;

public class mq {

    // MQ Queue Manager Configuration.
    private static String queueManager = "";
    private static String queueName = "";
    private static String channel = "";
    private static int port = 0;

    // Cognito JWKS Configuration.
    private static String username = "";
    private static String password = "";
    private static String clientId = "";
    private static String region = "";

    /*
        Purpose: This Java application demonstrates how to send and receive messages to/from an IBM MQ queue using
        AWS Cognito for JWKS-based authentication.

        Dependencies:
        - Jakarta Messaging API (JMS) for MQ interaction.
        - IBM MQ Java client libraries.
        - AWS SDK for Java (specifically, HttpClient for Cognito token retrieval).

        Configuration:
        Environment variables for MQ connection details (queue manager, queue name, channel, port) and Cognito configuration details
        (username, password, client ID, region) are loaded from the `env.json` file.

        Flow:
        - The application loads environment variables from env.json.
        - Retrives an access token from AWS Cognito using the provided credentials and the HTTP package.
        - Connects to the MQ queue and sends/receives messages specified by the arguments.
        - After successful message exchange, it prints a success message and exits.
        */

    public static void main(String[] args) {
        System.out.println("\nJava Jakarta App to send message to MQ with AWS Cognito providing JWKS Authentication.\n");

        // Load environment variables from env.json.
        getEnvVars();

        // Obtain access token from Cognito.
        String accessToken = getCognitoAccessToken();

        System.out.println("\nConnecting to MQ...");
        if (args.length > 0) {
            // Connect to MQ and use the command line arguments to send custom messages.
            connectAndTestMQ(accessToken, args);
        } else {
            // Use default message if no arguments are provided
            System.out.println("No arg messages found, using default test message.");
            String[] msgs = {"Hello from MQ using AWS Cognito JWKS Token Authtentication!"};
            // Connect to MQ using the defualt message.
            connectAndTestMQ(accessToken, msgs);
        }
        System.out.println("\nSuccessfully completed MQ Cognito JWKS validation test.");
        // Expected end of execution.
        System.exit(0);
    }

    // connectAndTestMQ, establishes a connection to an MQ queue using provided credentials and sends/receives messages.
    private static void connectAndTestMQ(String accessToken, String[] msgs){
        JMSContext context = null;
        Destination destination = null;
        JMSProducer producer = null;
        JMSConsumer consumer = null;

        try {
            System.out.println("Creating connection factory.");
            JmsFactoryFactory ff = JmsFactoryFactory.getInstance(WMQConstants.JAKARTA_WMQ_PROVIDER);
            JmsConnectionFactory cf = ff.createConnectionFactory();

            // Set MQ connection properties.
            cf.setIntProperty(WMQConstants.WMQ_CONNECTION_MODE, WMQConstants.WMQ_CM_CLIENT);
            cf.setIntProperty(WMQConstants.WMQ_PORT, port);
            cf.setStringProperty(WMQConstants.WMQ_CHANNEL, channel);
            cf.setStringProperty(WMQConstants.WMQ_QUEUE_MANAGER, queueManager);
            // Use the Cognito access token for authentication.
            cf.setStringProperty(WMQConstants.PASSWORD, accessToken);

            context = cf.createContext();
            destination = context.createQueue("queue:///" + queueName);

            System.out.println("Successfully connected to MQ Queue Manager and created queue.");
            System.out.println("\nPutting and Getting messages to the Queue...");

            for (String msg: msgs){
                producer = context.createProducer();
                consumer = context.createConsumer(destination);

                // Send a message
                TextMessage sendMessage = context.createTextMessage(msg);
                producer.send(destination, sendMessage);
                System.out.println("Sent message: " + sendMessage.getText());

                // Recieve a message
                TextMessage receivedMessage = (TextMessage) consumer.receive(10000);
                if (receivedMessage != null){
                    System.out.println("Received message: " + receivedMessage.getText());
                } else {
                    System.out.println("error: no message received?!");
                    System.exit(1);
                }
            }
        } catch (JMSException ex) {
            System.out.println("error JMSException occurred: " + ex.toString());
            System.exit(1);
        }
    }

    // getCognitoAccessToken, retrieves an access token from AWS Cognito using provided credentials.
    // It constructs a request, sends it to Cognito, and extracts the access token from the response.
    private static String getCognitoAccessToken(){
        System.out.println("\nGetting user access token from Amazon AWS Cognito...");

        String accessToken = "";
        String requestData = "{" +
            "  \"AuthParameters\": {" +
            "    \"USERNAME\": \"" + username + "\"," +
            "    \"PASSWORD\": \"" + password + "\"" +
            "  }," +
            "  \"AuthFlow\": \"USER_PASSWORD_AUTH\"," +
            "  \"ClientId\": \"" + clientId + "\"" +
            "}";
        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
            .uri(URI.create("https://cognito-idp." + region + ".amazonaws.com/"))
            .POST(BodyPublishers.ofString(requestData))
            .setHeader("X-Amz-Target", "AWSCognitoIdentityProviderService.InitiateAuth")
            .setHeader("Content-Type", "application/x-amz-json-1.1")
            .build();

        System.out.println("Cognito Request: " + request + requestData);

        try{
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            String JWT = response.body();

            Pattern accessTokenPattern = Pattern.compile("\"AccessToken\"\s*:\s*\"([^,]*)\",");
            Matcher accessTokenMatcher = accessTokenPattern.matcher(JWT);

            if (accessTokenMatcher.find()) {
                accessToken = accessTokenMatcher.group(1);
                System.out.println("Cognito Access Token: " + accessToken);
            } else {
                System.out.println("error: Could not parse JWT token response");
                System.exit(1);
            }
        } catch(IOException ex){
            System.out.println(ex.toString());
            System.exit(1);
        } catch (InterruptedException e) {
            System.out.println(e.toString());
            System.exit(1);
        }
        return accessToken;
    }

    // getEnvVars, loads environment variables from env.json.
    private static void getEnvVars(){
        String filename = "env.json";
        System.out.printf("Loading environment variables from %s...\n", filename);
        String envs = readFile(filename);

        System.out.println("MQ Config:");
        queueManager = getVar(envs, "QMGR");
        queueName = getVar(envs, "QUEUE_NAME");
        channel = getVar(envs, "CHANNEL");

        String portStr = getVar(envs, "PORT");
        try{
            port = Integer.parseInt(portStr);
        } catch (NumberFormatException e) {
            System.out.printf("error: port value provided [%s] is not a number: %s", portStr, e.toString());
            System.exit(1);
        }

        System.out.println("\nCognito Config:");
        username = getVar(envs, "USERNAME");
        password = getVar(envs, "PASSWORD");
        clientId = getVar(envs, "CLIENT_ID");
        region = getVar(envs, "REGION");
    }

    // getVar, retrieves the value of a specific environment variable from a JSON-formatted string.
    // If the variable is found, it returns its value; otherwise, it prints an error message and exits the program.
    private static String getVar(String envs, String name){
        Pattern pattern = Pattern.compile("\"" + name + "\"\s*:\s*\"([^,]*)\"");
        Matcher matcher = pattern.matcher(envs);

        if (matcher.find()) {
            String var = matcher.group(1);

            if (var == ""){
                System.out.println("error: value of " + name + ", is a empty string.");
                System.exit(1);
            }

            System.out.printf(" - %s: \t%s\n", name, var);
            return var;
        } else {
            System.out.println("error: Could not find env var " + name + ", check JSON format is correct and variable isn't missing.");
            System.exit(1);
        }
        return "";
    }

    // Readfile, reads contents of filename and retuns the file content as a String.
    // If error occurs, prints the eror and exits the program.
    private static String readFile(String filename) {
		try {
			return new String(Files.readAllBytes(Paths.get(filename)));
		} catch (IOException e) {
            System.out.printf("error: failed to read file [%s]: %s", filename, e.toString());
            System.exit(1);
            return null;
		}
	}
}
EOF

echo "mq.java created with the following content:"
cat mq.java

cd ..
echo "Setup of MQ Jakarta app is complete."