from typing import Dict, Union

from loguru import logger
from openai import OpenAI

from sawaw.method import sawaw_analyze_method
from sawaw.entry import SAWAWEntry
from sawaw.sentiments import SentimentResult

client = OpenAI(api_key="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")


def query(
    comment: str, aspect_word_or_words: Union[list, str]
) -> Union[Dict[str, SentimentResult], SentimentResult]:
    """
    Query the GPT-3 model for sentiment analysis.

    :param comment: The comment to be analyzed.
    :param aspect_word_or_words: The aspect word(s) to be analyzed. If a list is provided, the model will return a dictionary of aspect word to sentiment result. If a string is provided, the model will return a single sentiment result.
    :return: A dictionary of aspect word to sentiment result if a list is provided, or a single sentiment result if a string is provided.
    """

    def to_sentiment_result(result: str) -> SentimentResult:
        if result == "positive":
            return SentimentResult.POSITIVE
        elif result == "neutral":
            return SentimentResult.NEUTRAL
        elif result == "negative":
            return SentimentResult.NEGATIVE
        elif result == "none":
            return SentimentResult.NONE
        else:
            return SentimentResult.UNDEFINED

    if isinstance(aspect_word_or_words, list):
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "role": "system",
                    "content": "You will be given a comment and several aspect words as input. For each aspect word, please respond with exactly one of the following words: 'positive', 'negative', 'neutral', or 'none'. Do not respond with any other words.\n\ne.g. Comment: I love this pair of shoes, but I think the food nearby is poor :(\n\nAspect Word: shoes, food\n\nResponse: shoes: positive\nfood: negative\n\nPlease strictly follow the format above.",
                },
                {
                    "role": "user",
                    "content": "Comment:\n\n{}\n\nAspect Words:\n\n{}".format(
                        comment, ", ".join(aspect_word_or_words)
                    ),
                },
            ],
        )
        result = response.choices[0].message.content
        logger.debug("Result: {}".format(result))
        results = result.split("\n")
        aspect_word_to_result = {}
        for result in results:
            try:
                aspect_word, sentiment = result.split(":")
                aspect_word_to_result[aspect_word.strip()] = to_sentiment_result(
                    sentiment.strip().lower()
                )
            except:
                logger.warning("Invalid result: {}".format(result))
        for aspect_word in aspect_word_or_words:
            if aspect_word not in aspect_word_to_result:
                aspect_word_to_result[aspect_word] = SentimentResult.UNDEFINED
                logger.warning("Aspect word '{}' not found in result".format(aspect_word))
        return aspect_word_to_result
    elif isinstance(aspect_word_or_words, str):
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "role": "system",
                    "content": "You will be given a comment and an aspect word as input. Please respond with exactly one of the following words: 'positive', 'negative', 'neutral', or 'none'. Do not respond with any other words.",
                },
                {
                    "role": "user",
                    "content": "Comment:\n\n{} Aspect Word:\n\n{}".format(
                        comment, aspect_word_or_words
                    ),
                },
            ],
        )
        result = response.choices[0].message.content.lower()
        return to_sentiment_result(result)


@sawaw_analyze_method
def analyze(entry: SAWAWEntry) -> None:
    """
    Analyze the sentiment of an entry using the GPT-3 model. Modifies the entry in-place.

    :param entry: The entry to be analyzed.
    :return: The entry with the sentiment_results field populated.
    """
    results = query(entry.comment, entry.aspect_words)
    sentiment_results = []
    for aspect_word in entry.aspect_words:
        sentiment_results.append(results[aspect_word])
    entry.sentiment_results = sentiment_results