"""Store ham and spam for labeling and training."""
import json
import time
from threading import Thread

from google.cloud import storage
from vyper import v

from app import logger

log = logger.logger
bucket = None  # pylint: disable=invalid-name
_BUCKET_KEY = "gcs_bucket"
_stores = {}  # dict containing the DataStore store objects for each spammable type

if v.is_set(_BUCKET_KEY):
    bucket_str = v.get_string(_BUCKET_KEY)
    _storage_client = storage.Client()
    try:
        bucket = _storage_client.get_bucket(bucket_str)
        log.debug("GCS storage enabled")
    except Exception as exp:  # pylint: disable=broad-except
        log.error(f"Failed to initialize GCS bucket: {exp}")
else:
    log.debug("GCS storage disabled")


class DataStore:  # pylint: disable=too-few-public-methods
    """The data store is used to capture ham and spam and write to GCS for labeling and training.
    All issues classified as spam are saved and ham is saved in a balanced manner. Ham storage
    is biased to ensure that spammables that are closest to the SPAM_HAM_THRESHOLD are saved more
    often than other ham.

    Args:
        spammable_type (str): The type of spammable (i.e. issue, snippet).
    """

    MAX_HAM_SIZE = 50
    SPAM_HAM_THRESHOLD = 0.5
    QUESTIONABLE_HAM_BIAS = 3

    def __init__(self, spammable_type: str):
        self.spammable_type = spammable_type.lower()
        self.gcs_file_path = f"spamcheck-unlabeled/{spammable_type}"
        self._questionable_ham = {}
        self._random_ham = {}
        self._questionable_count = 1
        self._labeled_count = 0
        self._ham_spam_count = 0
        self._monitor_objects()

    def save(self, spammable: dict, confidence: float) -> None:
        """Save a spammable for future labeling and training.

        Args:
            spammable (dict): The dictionary representation of a spammable object.
            confidence (float): The confidence value returned from ML inference.
        """
        spammable["type"] = self.spammable_type
        spammable["ml_inference_score"] = confidence

        if confidence < self.SPAM_HAM_THRESHOLD:
            self._add(spammable, confidence)
        else:
            _write_to_gcs(self.gcs_file_path, spammable)
            self._ham_spam_count -= 1

        self._save_ham()

    def _save_ham(self) -> None:
        """Save ham messages until the saved ham is equal to the number of saved spam."""
        while self._ham_spam_count + self._labeled_count < 0:
            ham = self._get()
            if not ham:
                break
            _write_to_gcs(self.gcs_file_path, ham)

            # If _ham_spam_count is less than 0 then we are imbalanced in the unlabeled data we
            # have saved so balance that first.  Otherwise, we are here because of an imbalance
            # in the labeled data so decrement that counter instead.
            if self._ham_spam_count < 0:
                self._ham_spam_count += 1
            else:
                self._labeled_count += 1

    def _add(self, spammable: dict, confidence: float) -> None:
        # First fill up the "questionable ham". This dictionary contains spammables
        # That are closest to the threshold of being classified as spam.
        if len(self._questionable_ham) == 0 or confidence > min(self._questionable_ham):
            self._questionable_ham[confidence] = spammable
            # If the "questionsable ham" is full then remove the lowest confidence spammable
            if len(self._questionable_ham) > self.MAX_HAM_SIZE:
                self._questionable_ham.pop(min(self._questionable_ham))

        # If the "random ham" is not full then add the spammable to that dictionary
        elif len(self._random_ham) < self.MAX_HAM_SIZE:
            self._random_ham[confidence] = spammable

    def _get(self) -> dict:
        ham = None

        # Only get random ham if the modulo of the questionable count is 0.
        # Doing this will create a bias to return the ham messages that are closest
        # to the threshold of marking a spammable as spam versus ham.
        if (
            self._questionable_count % self.QUESTIONABLE_HAM_BIAS == 0
            and len(self._random_ham) > 0
        ):
            self._questionable_count = 1
            # popitem returns a key and value tuple but we have no use for the key
            _, ham = self._random_ham.popitem()

        elif len(self._questionable_ham) > 0:
            self._questionable_count += 1
            ham = self._questionable_ham.pop(max(self._questionable_ham))

        return ham

    def _monitor_objects(self):
        def update_counts():
            prefix = self.gcs_file_path.replace('unlabeled', 'labeled')

            while True:
                count = 0

                log.debug("counting labeled objects for {self.spammable_type}")
                for blob in bucket.list_blobs(prefix=prefix):
                    if '/spam/' in blob.name:
                        count -= 1
                    else:
                        count += 1

                log.debug(f"setting labeled count for {self.spammable_type}: {count}")
                self._labeled_count = count
                time.sleep(900)

        Thread(target=update_counts, daemon=True).start()


def save(spammable_type: str, spammable: dict, confidence: float):
    """Save a spammable for future labeling and training.

    Args:
        spammable_type (str): The type of spammable (i.e. issue, snippet).
        spammable (dict): The dictionary representation of the spammable object.
        confidence (float): The value confidence value returned from ML inference.
    """
    if bucket is None:
        return

    if spammable_type not in _stores:
        _stores[spammable_type] = DataStore(spammable_type)

    func = _stores[spammable_type].save
    thread = Thread(target=func, args=((spammable, confidence)))
    thread.start()


def _write_to_gcs(upload_path: str, data_dict: dict) -> None:
    """Convert a dict to JSON and write to GCS.

    Args:
        upload_path (str): The bucket path to write the blob.
        spammable_dict (dict): The dictionary to save.
    """
    json_str = json.dumps(data_dict)
    filename = data_dict["correlation_id"]
    try:
        blob = bucket.blob(f"{upload_path}/{filename}.json")
        blob.upload_from_string(json_str)
    # pylint: disable=broad-except
    except Exception as ex:
        log.error(f"failed to write to GCS: {ex}")
