import string

from loguru import logger
from pydantic import BaseModel, ConfigDict

from huesoporro.libs.db import MarkovDatabase as MarkovDB
from huesoporro.libs.tokenizer import detokenize, tokenize


class SentenceGeneratorSvc(BaseModel):
    db: MarkovDB
    min_sentence_length: int = 2
    key_length: int = 2
    max_sentence_length: int = 25
    sentence_separator: str = " "
    model_config = ConfigDict(arbitrary_types_allowed=True)

    @staticmethod
    def get_sentence_length(sentences: list[list[str]]) -> int:
        """Given a list of tokens representing a sentence, return the number of words in there.

        Args:
            sentences (List[List[str]]): List of lists of tokens that make up a sentence,
                where a token is a word or punctuation. For example:
                [['Hello', ',', 'you', "'re", 'Tom', '!'], ['Yes', ',', 'I', 'am', '.']]
                This would return 6.

        Returns:
            int: The number of words in the sentence.
        """
        count = 0
        for sentence in sentences:
            for token in sentence:
                if token not in string.punctuation and token[0] != "'":
                    count += 1
        return count

    def generate(self, params: list[str] | None = None) -> tuple[str, bool]:  # noqa: C901, PLR0912
        """Given an input sentence, generate the remainder of the sentence using the learned data.

        Args:
            params (list[str]): A list of words to use as an input to use as the start of generating.

        Returns:
            tuple[str, bool]: A tuple of a sentence as the first value, and a boolean indicating
                whether the generation succeeded as the second value.
        """
        params = params or []

        # List of sentences that will be generated. In some cases, multiple sentences will be generated,
        # e.g. when the first sentence has fewer words than self.min_sentence_length.
        sentences: list[list | list[str]] = [[]]

        # Get the starting key and starting sentence.
        # If there is more than 1 param, get the last 2 as the key.
        # Note that self.key_length is fixed to 2 in this implementation
        if len(params) > 1:
            key = params[-self.key_length :]
            # Copy the entire params for the sentence
            sentences[0] = params.copy()

        elif len(params) == 1:
            # First we try to find if this word was once used as the first word in a sentence:
            key = self.db.get_next_single_start(params[0])  # type: ignore[assignment]
            if key is None:
                # If this failed, we try to find the next word in the grammar as a whole
                key = self.db.get_next_single_initial(0, params[0])
                if key is None:
                    # Return a message that this word hasn't been learned yet
                    return f'I haven\'t extracted "{params[0]}" from chat yet.', False
            # Copy this for the sentence
            sentences[0] = key.copy()

        else:  # if there are no params
            # Get starting key
            key = self.db.get_start()
            if key:
                # Copy this for the sentence
                sentences[0] = key.copy()
            else:
                # If nothing's ever been said
                return "There is not enough learned information yet.", False

        # Counter to prevent infinite loops (i.e. constantly generating <END> while below the
        # minimum number of words to generate)
        i = 0
        while (
            self.get_sentence_length(sentences) < self.max_sentence_length
            and i < self.max_sentence_length * 2
        ):
            # Use key to get next word
            if i == 0:
                # Prevent fetching <END> on the first word
                word = self.db.get_next_initial(i, key)
            else:
                word = self.db.get_next(i, key)

            i += 1

            if word == "<END>" or word is None:
                # Break, unless we are before the min_sentence_length
                if i < self.min_sentence_length:
                    key = self.db.get_start()
                    # Ensure that the key can be generated. Otherwise, we still stop.
                    if key:
                        # Start a new sentence
                        sentences.append([])
                        for entry in key:
                            sentences[-1].append(entry)
                        continue
                break

            # Otherwise add the word
            sentences[-1].append(word)

            # Shift the key so on the next iteration it gets the next item
            key.pop(0)
            key.append(word)

        # If there were params, but the sentence resulting is identical to the params
        # Then the params did not result in an actual sentence
        # If so, restart without params
        if len(params) > 0 and params == sentences[0]:
            return "I haven't learned what to do with \"" + detokenize(
                params[-self.key_length :],
            ) + '" yet.', False

        return self.sentence_separator.join(
            detokenize(sentence) for sentence in sentences
        ), True

    async def run(
        self,
        sentence: str | None = None,
    ) -> str | None:
        split_sentence = tokenize(sentence) if sentence else None

        logger.info(f"Generating sentence from {split_sentence}")
        generated_sentence, success = self.generate(split_sentence)
        logger.info(f"Generated sentence: {generated_sentence}")
        if not success:
            return None
        return generated_sentence
