package Coordinator_Worker;

import Zookeeper.ServiceRegistry;

import java.io.IOException;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;

public class Coordinator {

    // Method to start the coordinator and coordinate the processing
    public static TreeMap<String, Double> startCoordinator(String dataString) {
        // Data structures to store results
        ConcurrentLinkedQueue<Document> workersResult = new ConcurrentLinkedQueue<>();
        TreeMap<String, Double> IDF = new TreeMap<>();

        try {
            // Split data string into terms
            String[] parts = dataString.split("\\s");
            List<Term> terms = new ArrayList<>();
            for (String part : parts) {
                terms.add(new Term(part));
            }

            // Response data structure
            TreeMap<String, Double> response = new TreeMap<>();

            // Get file names in the specified directory
            List<String> fileNames = getFileNamesInDirectory("src/main/resources/files");

            // Calculate the number of files each worker should handle
            int eachWorkerFiles = fileNames.size() / ServiceRegistry.getWorkers().size();
            int mode = fileNames.size() % ServiceRegistry.getWorkers().size();
            int index = 0;

            // Thread pool to manage worker threads
            ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

            // Check if workers are available
            if (ServiceRegistry.getWorkers() != null) {
                for (Worker worker : ServiceRegistry.getWorkers()) {
                    // Prepare the list of files to be processed by the worker
                    List<String> workerFilesList = new ArrayList<>();
                    for (int i = 0; i < eachWorkerFiles; i++) {
                        workerFilesList.add(fileNames.get(index));
                        index++;

                        // Handle the remaining files if any and the mode is not zero
                        if (mode != 0) {
                            if (index != fileNames.size()) {
                                workerFilesList.add(fileNames.get(index));
                                index++;
                            }
                        }
                    }

                    // Connect to the worker
                    Socket workerSocket = new Socket(worker.getIpAddress(),worker.getPort() );

                    // Create a CoordinatorHandler instance for each worker
                    Runnable task = () -> processRequest(workerSocket, workerFilesList, dataString, workersResult);

                    // Start the CoordinatorHandler thread
                    executorService.submit(task);
                }

                // Wait for all worker threads to finish
                executorService.awaitTermination(5, TimeUnit.SECONDS);

                // Process results to calculate Term Frequency (TF)
                for (int i = 0; i < terms.size(); i++) {
                    for (Document document : workersResult) {
                        Map<String, Double> tf = document.getTf();
                        if (tf.get(terms.get(i).getTermName()) > 0) {
                            terms.get(i).increase();
                        }
                    }
                }

                // Set terms and calculate IDF for each document
                for (Document document : workersResult) {
                    document.setTerms(terms);
                    document.CalculateIDF(fileNames.size());
                    IDF.put(document.getName(), document.getIDF());
                }

                // Shutdown the executor service
                executorService.shutdown();
            }

        } catch (UnknownHostException e) {
            // Handle unknown host exception
            throw new RuntimeException(e);
        } catch (IOException e) {
            // Handle IO exception
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            // Handle interrupted exception
            throw new RuntimeException(e);
        } finally {
            // Return the calculated IDF values
            return IDF;
        }
    }

    // Method to get file names in the specified directory
    public static List<String> getFileNamesInDirectory(String directoryPath) {
        try {
            // Create a Path object for the specified directory
            Path directory = Paths.get(directoryPath);

            // Use try-with-resources to open a DirectoryStream and iterate over the files
            try (DirectoryStream<Path> directoryStream = Files.newDirectoryStream(directory)) {
                // Create a list to store file names
                List<String> fileNames = new ArrayList<>();

                // Iterate over the files in the directory
                for (Path path : directoryStream) {
                    // Get the file name and add it to the list
                    fileNames.add(path.getFileName().toString());
                }

                // Return the list of file names
                return fileNames;
            }
        } catch (IOException e) {
            // Handle exceptions, e.g., directory not found
            e.printStackTrace();
            return Collections.emptyList();
        }
    }

    // Method to process requests from workers
    private static void processRequest(Socket clientSocket, List<String> fileNames,
                                       String dataString, ConcurrentLinkedQueue<Document> documents) {
        try {
            // Send data to worker
            Communication.sendData(clientSocket, dataString, fileNames);

            // Receive results from worker
            ConcurrentLinkedQueue<Document> d = Communication.receiveResult(clientSocket);
            documents.addAll(d);

            // Close the client socket
            clientSocket.close();

        } catch (IOException e) {
            // Handle IO exception
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            // Handle class not found exception
            throw new RuntimeException(e);
        }
    }
}
