# pragma: no cover
from loguru import logger
from nltk.tokenize import sent_tokenize
from pydantic import BaseModel, ConfigDict

from huesoporro.libs.db import MarkovDatabase
from huesoporro.libs.tokenizer import tokenize


class SentenceStorerSvc(BaseModel):
    db: MarkovDatabase
    key_length: int = 2
    end_tag: str = "<END>"

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def store_sentence(self, message: str):
        logger.info(f"Processing {message} in order to store it")
        stripped_message = message.strip()
        try:
            sentences = sent_tokenize(stripped_message)
        except LookupError:
            logger.debug("Downloading required punkt resource...")
            import nltk

            nltk.download("punkt")
            nltk.download("punkt_tab")
            logger.debug("Downloaded required punkt resource.")
            sentences = sent_tokenize(stripped_message)

        for sentence in sentences:
            words = tokenize(sentence)
            # Double spaces will lead to invalid rules. We remove empty words here
            if "" in words:
                words = [word for word in words if word]

            # If the sentence is too short, ignore it and move on to the next.
            if len(words) <= self.key_length:
                continue

            # Add a new starting point for a sentence to the <START>
            words = [words[x] for x in range(self.key_length)]
            logger.debug(f"Adding {words} to start queue")
            self.db.add_start_queue(words)

            # Create Key variable which will be used as a key in the Dictionary for the grammar
            key: list[str] = []
            for word in words:
                # Set up key for first use
                if len(key) < self.key_length:
                    key.append(word)
                    continue
                logger.debug(f"Adding {key}[{word}] to rule queue")
                self.db.add_rule_queue([*key, word])

                # Remove the first word, and add the current word,
                # so that the key is correct for the next word.
                key.pop(0)
                key.append(word)
            logger.debug(f"Adding {key} to rule queue")
            # Add <END> at the end of the sentence
            self.db.add_rule_queue([*key, self.end_tag])

    async def run(self, sentence: str):
        return self.store_sentence(sentence)
