import uuid
import threading
import math
from collections import Counter
from collections.abc import Generator
from abc import ABC
from typing import Any

from icefarm.utils import MappedQueues

class Evaluation:
    """
    Circuit evaluation to be sent to the worker. Note that providing
    multiple serials to one Evaluation will result in faster evaluations than
    providing multiple Evaluations of the same bitstream.
    """
    def __init__(self, serials: set[str]):
        if not serials:
            raise Exception("Serials cannot be empty")

        self.serials = frozenset(serials)
        self.id = str(uuid.uuid4())

    def __eq__(self, other):
        return self.id  == other.id

    def __hash__(self):
        return hash(self.id)

class EvaluationBundle:
    """
    Bundles Evaluations into efficient batches that can be sent gradually rather than
    all at once.
    """
    def __init__(self, evaluations: list[Evaluation], batch_size):
        self.queue = MappedQueues()
        for evaluation in evaluations:
            self.queue.append(evaluation.serials, evaluation)

        self.evaluation_lookup = {evaluation.id : evaluation for evaluation in evaluations}

        self.batch_size = batch_size

        self.empty = False
        self.id = str(uuid.uuid4())

    def __next__(self) -> dict[set[str], list[Evaluation]]:
        """
        Produces the next batch. Batches are produced so that each Evaluation with multiple serials are sent in the same batch
        and the maximum evaluations are sent without exceeding the batch size amount of circuits for any serial.
        """
        serial_amounts = Counter()
        batch = {}

        search_order = sorted(self.queue.keys(), key=len)
        for serials in search_order:
            slots = min(self.batch_size - serial_amounts[serial] for serial in serials)
            if not slots:
                continue

            commands = self.queue.pop(serials, slots)
            batch[serials] = commands
            for serial in serials:
                serial_amounts[serial] += len(commands)

        if not batch:
            self.empty = True
            raise StopIteration

        return batch

    def __iter__(self):
        return self

class AbstractBatchFactory(ABC):
    """
    Produces batches for client consumption.
    """
    def __init__(self, bundle: EvaluationBundle):
        super().__init__()
        self.bundle = bundle
        self.results: list[tuple[Evaluation, dict]] = []
        self.result_cv = threading.Condition()

        self.awaiting_results: dict[str, set] = {}

    def processResult(self, serial: str, evaluation_id: str, result: Any):
        with self.result_cv:
            if serial not in self.awaiting_results:
                return

            if evaluation_id not in self.awaiting_results[serial]:
                return

            self.awaiting_results[serial] -= {evaluation_id}
            self.results.append((serial, evaluation_id, result))
            self.result_cv.notify_all()

    def getResults(self) -> Generator[tuple[str, Evaluation, Any]]:
        """Produces results as they received by processResult."""
        while True:
            with self.result_cv:
                if not self.results:
                    self.result_cv.wait_for(lambda : self.results)

                for serial, evaluation_id, result in self.results:
                    evaluation = self.bundle.evaluation_lookup.get(evaluation_id)
                    yield serial, evaluation, result

                self.results = []

                if self.bundle.empty and not any(self.awaiting_results.values()):
                    return

    def _addBatch(self, batch: dict[set[str], list[Evaluation]]):
        """Adds batch to awaiting results."""
        for serials, evaluations in batch.items():
            for serial in serials:
                if serial not in self.awaiting_results:
                    self.awaiting_results[serial] = set()

                self.awaiting_results[serial] |= set(evaluation.id for evaluation in evaluations)

    def getBatches(self) -> Generator[dict[set[str], list[Evaluation]]]:
        """Produces batches for client to consume."""
        for batch in self.bundle:
            with self.result_cv:
                if not self._readyForBatch():
                    self.result_cv.wait_for(self._readyForBatch)

                self._addBatch(batch)
                yield batch

    def _readyForBatch(self) -> bool:
        """Whether the next batch should be made available for consumption."""

class QuickBatchFactory(AbstractBatchFactory):
    """
    Sends all batches to workers without any delay.
    """
    def _readyForBatch(self):
        return True

class PatientBatchFactory(AbstractBatchFactory):
    """
    Waits until all workers have finished processing a batch
    before sending a new one.
    """
    def _readyForBatch(self):
        return not any(self.awaiting_results.values())

class BalancedBatchFactory(AbstractBatchFactory):
    """
    Starts by sending an initial amount of batches, then
    waits until a threshold of results have been received
    before sending more batches.
    """
    def __init__(self, bundle, target_batches=4):
        super().__init__(bundle)
        self.target_batches = target_batches

    def _readyForBatch(self):
        if not any(self.awaiting_results.values()):
            return True

        highest_queued = max(map(len, self.awaiting_results.values()))
        batches = math.ceil(highest_queued / self.bundle.batch_size)
        return bool(self.target_batches - batches)
