package com.ease.gsms.server.services.dispatch;

import com.ease.gsms.server.model.Dispatch;
import com.ease.gsms.server.model.DispatchStatus;
import com.ease.gsms.server.model.Message;
import org.slf4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Controller;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

@Controller
public class WorkAllocator {

    @Autowired
    Logger logger;

    private Map<String, MessageSender> workerContainer;

    private Map<String, Map<Long, Dispatch>> workContainer;

    @Async
    void doAllocations(Set<String> usersWithChanges) {
        for (String user : usersWithChanges) {
            doAllocations(user);
        }
    }

    private void doAllocations(String user) {

        Map<String, Allocation> newAllocations = computeAllocations(user);

        Set<MessageSender> referencedWorkers = newAllocations.values().stream().map(Allocation::getWorker).collect(Collectors.toSet());

        // 1 - deallocate unreferenced workers
        Set<MessageSender> workersToDeallocate = workerContainer.values().stream().filter(worker -> user.equals(worker.getUser())).collect(Collectors.toSet());

        workersToDeallocate.removeAll(referencedWorkers);

        for (MessageSender worker : workersToDeallocate) {

            worker.deallocate();

        }

        // 2 - allocate new
        newAllocations
                .entrySet()
                .forEach(entry -> {
                    workerContainer.get(entry.getKey()).allocate(entry.getValue());
                });

    }

    private synchronized Map<String, Allocation> computeAllocations(String user) {

        Map<Long, Dispatch> dispatchesById = workContainer.getOrDefault(user, Collections.emptyMap());
        Set<MessageSender> workers = workerContainer.values().stream().filter(worker -> user.equals(worker.getUser())).collect(Collectors.toSet());

        Map<String, MessageSender> workersBySimId = workers.stream().collect(Collectors.toMap(MessageSender::getId, Function.identity()));
        Map<String, Allocation> allocationsBySimId = workersBySimId.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> new Allocation(entry.getValue())));

        // Step 1 - select only eligible dispatches

        Set<Dispatch> dispatches = dispatchesById.values().stream().filter(dispatch -> {
            // TODO: implement checks on allowed time intervals
            switch (dispatch.getStatus()) {
                case COMPLETED:
                case CANCELLED:
                case PAUSED:
                case SCHEDULED_PAUSE:
                    return false; // we don't allocate dispatches in these states
            }
            return true;
        }).collect(Collectors.toSet());

        // Step 2 - allocate dispatches with predefined senders

        for (Dispatch dispatch : dispatches) {
            String iccid = dispatch.getFromICCID();
            if (iccid != null && allocationsBySimId.containsKey(iccid)) { // has predefined sender and there is a worker that can operate on it
                allocationsBySimId.get(iccid).addTask(
                        new Task(
                                dispatch,
                                true,
                                dispatch.unsent().collect(Collectors.toList())
                        )
                );
            }
        }

        // Step 3 - distribute random allocations to spread work as evenly as possible in bins

        int[] bins = allocationsBySimId.values().stream().mapToInt(Allocation::getWeight).toArray();

        int[] unsentMessageCounts = dispatches.stream().mapToInt(dispatch -> dispatch.getFromICCID() != null ? -1 : dispatch.getUnsentMessageCount()).toArray();

        int totalUnallocatedMessageCount = 0;

        int i = 0;

        for (Dispatch dispatch : dispatches) {
            if (dispatch.getFromICCID() == null) {
                totalUnallocatedMessageCount += unsentMessageCounts[i];
            }
            i++;
        }

        int[] addToBin = BinBalancer.balanceDelta(totalUnallocatedMessageCount, bins); // how many additional messages should be provided to each worker

        if (!allocationsBySimId.isEmpty()) {

            Iterator<Allocation> allocationIterator = allocationsBySimId.values().iterator();

            Allocation currentAllocation = allocationIterator.next();

            i = 0;

            for (Dispatch dispatch : dispatches) {

                if (dispatch.getFromICCID() == null) {

                    List<Message> currentMessageList = new LinkedList<>();

                    for (Message message : dispatch.getMessages()) {

                        if (message.isUnsent()) {

                            if (addToBin[i]-- > 0) {

                                currentMessageList.add(message);

                            } else { // filled it up, move to next allocation

                                if (!currentMessageList.isEmpty()) {

                                    currentAllocation.addTask(
                                            new Task(dispatch, false, currentMessageList)
                                    );

                                    currentMessageList = new LinkedList<>();

                                }

                                currentMessageList.add(message);

                                currentAllocation = allocationIterator.hasNext() ? allocationIterator.next() : null;

                                i++;

                            }

                        }

                    }

                    if (!currentMessageList.isEmpty()) {

                        currentAllocation.addTask(
                                new Task(dispatch, false, currentMessageList)
                        );

                    }

                }

            }

        }

        allocationsBySimId = allocationsBySimId.entrySet().stream().filter(entry -> entry.getValue().hasNext()).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

        if (!allocationsBySimId.isEmpty()) {
            logger.info("Allocations for {}", user);
            for (Allocation allocation : allocationsBySimId.values()) {
                logger.info("{}: {} tasks, {} messages", allocation.getUuid(), allocation.getTasks().size(), allocation.getWeight());
                for (Task task : allocation.getTasks()) {
                    logger.info("{}: task with {} messages, depleted: {}", allocation.getUuid(), task.getWeight(), !task.hasNext());
                }
            }
        }

        return allocationsBySimId;

    }

    void setWorkerContainer(Map<String, MessageSender> workerContainer) {
        this.workerContainer = workerContainer;
    }

    void setWorkContainer(Map<String, Map<Long, Dispatch>> workContainer) {

        this.workContainer = workContainer;
    }

    public Map<String, MessageSender> getWorkerContainer() {
        return workerContainer;
    }

    public Map<String, Map<Long, Dispatch>> getWorkContainer() {
        return workContainer;
    }
}
