from sawaw.sentiments import SentimentResult
from dataclasses import dataclass
from typing import List
import colorama

color_mapping = {
    SentimentResult.POSITIVE: colorama.Fore.GREEN,
    SentimentResult.NEUTRAL: colorama.Fore.YELLOW,
    SentimentResult.NEGATIVE: colorama.Fore.RED,
    SentimentResult.NONE: colorama.Fore.LIGHTWHITE_EX,
    SentimentResult.UNDEFINED: colorama.Fore.MAGENTA,
}


@dataclass
class SAWAWEntry:
    comment: str
    aspect_words: List[str]
    sentiment_results: List[SentimentResult] = None

    def __post_init__(self):
        if self.sentiment_results is None:
            self.sentiment_results = [SentimentResult.UNDEFINED] * len(
                self.aspect_words
            )

    def __repr__(self):
        colored_comment = self.comment
        for aspect_word in self.aspect_words:
            colored_comment = colored_comment.replace(
                aspect_word,
                color_mapping[
                    self.sentiment_results[self.aspect_words.index(aspect_word)]
                ]
                + aspect_word
                + colorama.Style.RESET_ALL,
            )
        return colored_comment
    
    def __str__(self):
        return self.__repr__()

    @staticmethod
    def print_legends():
        legends = []
        for sentiment_result, color in color_mapping.items():
            legends.append(
                "{}{}{}".format(color, sentiment_result.name, colorama.Style.RESET_ALL)
            )
        print(" | ".join(legends))