import random
import logging
from velascobot2.database import Database
from velascobot2.utils import GraphicType, Response, ChatData

logger = logging.getLogger(__name__)


def _rewrite(text: str) -> list[str]:
    """
    Splits strings into lists of words delimited by whitespaces.
    Extra whitespaces are inserted around newline characters so they are
    included as their own Markov chain element, to avoid polluting with
    very similar words that would only differ in having a whitespace
    attached or not
    """
    words = text.replace("\n", " \n ").split(" ")
    return [word.strip(" ") for word in words if len(word.strip(" ")) > 0]


def _normalize_key(*words: str) -> str:
    """Returns a case insensitive n-gram key by joining them with whitespaces"""
    return " ".join(word.casefold() for word in words)


def _triplets(wordlist: list[str]):
    """Generates triplets of words from the given data string. So if our string
    were "What a lovely day", we'd generate (What, a, lovely) and then
    (a, lovely, day).
    """
    if len(wordlist) < 3:
        return

    for i in range(len(wordlist) - 2):
        yield (wordlist[i], wordlist[i + 1], wordlist[i + 2])


class Parrot:
    def __init__(self, bot_name: str):
        self.db = Database(bot_name)

    def parse_text(self, chat_id: str, text: str) -> None:
        logger.debug(f"{chat_id} - Processing text: {text}")
        words = [w for w in _rewrite(text) if w]
        words.append("")  # end of message
        words = ["", ""] + words  # start of message

        for key, word, count in self.parse_words(words):
            self.db.add_words(chat_id, key, word, count)

    def parse_words(self, words: list[str]):
        word_counts: dict[tuple[str, str], int] = {}

        for w1, w2, w3 in _triplets(words):
            key = _normalize_key(w1, w2)
            word_counts[(key, w3)] = word_counts.get((key, w3), 0) + 1

        for (key, word), count in word_counts.items():
            yield (key, word, count)

    def parse_graphic(
        self, chat_id: str, graph_type: GraphicType, graph_id: str
    ):
        logger.debug(f"{chat_id} - Parsing graphic {graph_type}: {graph_id}")
        assert graph_type in GraphicType
        self.db.add_graphic(chat_id, graph_type, graph_id)

    def give_text(self, chat_id: str) -> str:
        w1, w2 = "", ""
        w3 = "invalid"
        words = []
        while w3:
            key = _normalize_key(w1, w2)
            weighted_next_words = self.db.get_next_words(chat_id, key)
            w3 = self.choose_word(weighted_next_words)
            if w3 is None:
                raise Exception(f"Got a nonexistent word. Generated text so far: {' '.join(words)}. Dropping text")
            words.append(w3)
            w1, w2 = w2, w3
        return " ".join(words)

    def give_graphic(self, chat_id: str) -> tuple[GraphicType, str]:
        weighted_graphics = self.db.get_graphics(chat_id)
        return self.choose_graphic(weighted_graphics)

    def choose_word(self, weighted_words: list[tuple[str, int]]) -> str | None:
        if len(weighted_words) == 0:
            return None
        words, weights = zip(*weighted_words)
        return random.choices(words, weights=weights)[0]

    def choose_graphic(
        self, weighted_graphics: list[tuple[str, str, int]]
    ) -> tuple[GraphicType, str]:
        graphics = [
            (graph_type, graph_id)
            for graph_type, graph_id, _ in weighted_graphics
        ]
        weights = [count for _, _, count in weighted_graphics]
        chosen_type, chosen_id = random.choices(graphics, weights=weights)[0]
        return GraphicType(chosen_type), chosen_id

    def choose_response_type(self, chat_id: str) -> Response:
        graphics_count = self.db.count_graphics(chat_id)
        words_count = self.db.count_words(chat_id, _normalize_key("", ""))
        total = graphics_count + words_count

        if total > 0 and random.random() < graphics_count / total:
            return Response.GRAPHIC
        else:
            return Response.TEXT

    def get_all_chats(self) -> list[ChatData]:
        return self.db.get_all_chats()

    def get_chat_data(self, chat_id: str) -> ChatData | None:
        return self.db.get_chat_data(chat_id)

    def ensure_chat_exists(
        self, chat_id: str, chat_type: str, chat_name: str
    ) -> bool:
        return self.db.ensure_chat_exists(chat_id, chat_type, chat_name)

    def update_chat_name(self, chat_id: str, chat_name: str):
        self.db.update_chat_name(chat_id, chat_name)

    def update_chat_answer_chance(self, chat_id: str, answer_chance: float):
        self.db.update_chat_answer_chance(chat_id, answer_chance)

    def update_chat_message_period(self, chat_id: str, period: int):
        self.db.update_chat_message_period(chat_id, period)

    def update_chat_set_restricted(self, chat_id: str, restricted: bool):
        self.db.update_chat_set_restricted(chat_id, restricted)

    def update_chat_set_silenced(self, chat_id: str, silenced: bool):
        self.db.update_chat_set_silenced(chat_id, silenced)

    def get_processed_messages(self, chat_id: str) -> int:
        graph_count = self.db.get_chat_graph_count(chat_id)
        start_words = self.db.get_next_words(chat_id, " ")
        start_word_count = sum(weight for _, weight in start_words)
        return graph_count + start_word_count