import re
from functools import reduce

from app.configs.constants import TemperatureSettings, GPTModels
from app.helpers import count_tokens
from app.services.abc import ILLMService, ITrainingService


class TrainingService(ITrainingService):

    def __init__(self, llm: ILLMService):
        self._llm = llm

    async def fetch_tips(self, context: str, question: str, answer: str, correct_answer: str):
        messages = self._get_question_tips(question, answer, correct_answer, context)

        token_count = reduce(lambda count, item: count + count_tokens(item)['n_tokens'],
                             map(lambda x: x["content"], filter(lambda x: "content" in x, messages)), 0)

        response = await self._llm.prediction(
            GPTModels.GPT_3_5_TURBO,
            messages,
            None,
            TemperatureSettings.TIPS_TEMPERATURE,
            token_count=token_count
        )

        if isinstance(response, str):
            response = re.sub(r"^[a-zA-Z0-9_]+\:\s*", "", response)

        return response

    @staticmethod
    def _get_question_tips(question: str, answer: str, correct_answer: str, context: str = None):
        messages = [
            {
                "role": "user",
                "content": (
                    "You are a IELTS exam program that analyzes incorrect answers to questions and gives tips to "
                    "help students understand why it was a wrong answer and gives helpful insight for the future. "
                    "The tip should refer to the context and question."
                ),
            }
        ]

        if not (context is None or context == ""):
            messages.append({
                "role": "user",
                "content": f"This is the context for the question: {context}",
            })

        messages.extend([
            {
                "role": "user",
                "content": f"This is the question: {question}",
            },
            {
                "role": "user",
                "content": f"This is the answer: {answer}",
            },
            {
                "role": "user",
                "content": f"This is the correct answer: {correct_answer}",
            }
        ])

        return messages