# Needed to use PSC build
import exec_anaconda

exec_anaconda.exec_anaconda_or_die()

# Needed to import libraries in /lib folder
import sys, os
import time
import uuid

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "lib"))
from splunklib.searchcommands import dispatch, GeneratingCommand, Configuration, Option, Boolean

import setup_logging

logger = setup_logging.get_logger()


def is_meme(spl):
    from collections import Counter
    split_str = spl.split()
    if len(split_str) < 10:
        return False

    token_counts = Counter(split_str)
    for token in token_counts:
        if len(token) > 1 and token_counts[token] > len(split_str) // 4:
            return True

    return False


def postprocess(choices):
    not_meme = [s for s in choices if not is_meme(s)]
    spl_pred = not_meme[0] if len(not_meme) > 0 else choices[0]
    return spl_pred


@Configuration()
class SPLGenCommand(GeneratingCommand):
    """
    Generating command that yields predicted SPL queries based on a plain English description of a query.
    """

    # Required plain English description of a query
    prompt = Option(require=True)
    explain = Option(default=False, require=False, validate=Boolean())

    def generate(self):
        """
        Yields records of predicted SPL queries to the Splunk processing pipeline.
        """
        start_time = time.time()
        logger.info("Starting generate.")
        logger.info(f"UUID={str(uuid.uuid4())}")
        # Needed to reduce overhead in "Parsing search"
        from spl_gen.t5onnx.predict import Predictor


        predictor = Predictor()
        # Make predictions
        prompt = self.prompt.lower()
        if self.explain:
            predictions = predictor.predict(prompt)
            explain = predictions[0]
            yield self.gen_record(
                EXPLAIN=explain
            )
        else:
            predictions = predictor.predict(prompt)
            spl = postprocess(predictions)
            if "index=*" in spl or "sourcetype=*" in spl:
                self.write_warning("Using 'index=*' or 'sourcetype=*' can potentially be very expensive for your stack.")
            if "| join" in spl:
                self.write_warning("The ‘join’ operation is usually very expensive.")
            yield self.gen_record(
                SPL=spl,
            )

        end_time = time.time()
        logger.info(f"command=splgen, apply_time={round((end_time - start_time), 5)}")
        logger.info("Exiting generate.")


dispatch(command_class=SPLGenCommand)