import asyncio
import os
from typing import Dict, Optional
from uuid import uuid4

from ielts_be.configs.constants import GPTModels, TemperatureSettings, FilePaths
from ielts_be.helpers import TextHelper, ExercisesHelper, FileHelper
from ielts_be.repositories import IFileStorage
from ielts_be.services import ILLMService, IAIDetectorService


class GradeWriting:

    def __init__(self, llm: ILLMService, file_storage: IFileStorage, ai_detector: IAIDetectorService):
        self._llm = llm
        self._file_storage = file_storage
        self._ai_detector = ai_detector

    async def grade_writing_task(self, task: int, question: str, answer: str, attachment: Optional[str] = None):
        bare_minimum = 100 if task == 1 else 180

        if not TextHelper.has_words(answer):
            return self._zero_rating("The answer does not contain enough english words.")
        elif not TextHelper.has_x_words(answer, bare_minimum):
            return self._zero_rating("The answer is insufficient and too small to be graded.")
        else:
            template = self._get_writing_template()
            messages = [
                {
                    "role": "system",
                    "content": (
                        f'You are a helpful assistant designed to output JSON on this format: {template}'
                    )
                },
                {
                    "role": "user",
                    "content": (
                        f'Evaluate the given Writing Task {task} response based on the IELTS grading system, '
                        'ensuring a strict assessment that penalizes errors. Deduct points for deviations '
                        'from the task, and assign a score of 0 if the response fails to address the question. '
                        'Additionally, provide a detailed commentary highlighting both strengths and '
                        'weaknesses in the response. '
                        f'\n Question: "{question}" \n Answer: "{answer}"')
                }
            ]

            if task == 1:
                if attachment is None:
                    messages.append({
                        "role": "user",
                        "content": (
                            'Refer to the parts of the letter as: "Greeting Opener", "bullet 1", "bullet 2", '
                            '"bullet 3", "closer (restate the purpose of the letter)", "closing greeting"'
                        )
                    })
                else:
                    uuid = str(uuid4())
                    name = attachment.split('%2F')[-1].split('?')[0]
                    os.makedirs(f'./tmp/{uuid}/', exist_ok=True)
                    out_path = f'./tmp/{uuid}/{name}'
                    path = await self._file_storage.download_firebase_file(f'{FilePaths.WRITING_ATTACHMENTS}{name}', out_path)
                    b64_image = await FileHelper.encode_image(path)
                    messages.append(
                        {
                            "role": "user",
                            "content": [{
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/{name.split('.')[-1]};base64,{b64_image}"
                                    }
                            }]
                        }
                    )
                    FileHelper.remove_directory(f'./tmp/{uuid}')

            temperature = (
                TemperatureSettings.GRADING_TEMPERATURE
                if task == 1 else
                TemperatureSettings.GEN_QUESTION_TEMPERATURE
            )

            evaluation_promise = self._llm.prediction(
                GPTModels.GPT_4_O,
                messages,
                ["comment"],
                temperature
            )

            perfect_answer_minimum = 150 if task == 1 else 250
            perfect_answer_promise = self._get_perfect_answer(question, perfect_answer_minimum)
            fixed_text_promise = self._get_fixed_text(answer)
            ai_detection_promise = self._ai_detector.run_detection(answer)

            prediction_result, perfect_answer_result, fixed_text_result, ai_detection_result = await asyncio.gather(
                evaluation_promise,
                perfect_answer_promise,
                fixed_text_promise,
                ai_detection_promise
            )

            response = prediction_result
            response["perfect_answer"] = perfect_answer_result["perfect_answer"]
            response["overall"] = ExercisesHelper.fix_writing_overall(
                response["overall"],
                response["task_response"]
            )
            response['fixed_text'] = fixed_text_result

            if ai_detection_result is not None:
                response['ai_detection'] = ai_detection_result

            return response

    async def _get_fixed_text(self, text):
        messages = [
            {
                "role": "system",
                "content": (
                    'You are a helpful assistant designed to output JSON on this format: '
                    '{"fixed_text": "fixed test with no misspelling errors"}'
                )
            },
            {
                "role": "user",
                "content": (
                    'Fix the errors in the given text and put it in a JSON. '
                    f'Do not complete the answer, only replace what is wrong. \n The text: "{text}"'
                )
            }
        ]

        response = await self._llm.prediction(
            GPTModels.GPT_3_5_TURBO,
            messages,
            ["fixed_text"],
            0.2,
            False
        )
        return response["fixed_text"]

    async def _get_perfect_answer(self, question: str, size: int) -> Dict:
        messages = [
            {
                "role": "system",
                "content": (
                    'You are a helpful assistant designed to output JSON on this format: '
                    '{"perfect_answer": "perfect answer for the question"}'
                )
            },
            {
                "role": "user",
                "content": f'Write a perfect answer for this writing exercise of a IELTS exam. Question: {question}'

            },
            {
                "role": "user",
                "content": f'The answer must have at least {size} words'
            }
        ]
        return await self._llm.prediction(
            GPTModels.GPT_4_O,
            messages,
            ["perfect_answer"],
            TemperatureSettings.GEN_QUESTION_TEMPERATURE
        )

    @staticmethod
    def _zero_rating(comment: str):
        return {
            'comment': comment,
            'overall': 0,
            'task_response': {
                'Task Achievement': {
                    "grade": 0.0,
                    "comment": ""
                },
                'Coherence and Cohesion': {
                    "grade": 0.0,
                    "comment": ""
                },
                'Lexical Resource': {
                    "grade": 0.0,
                    "comment": ""
                },
                'Grammatical Range and Accuracy': {
                    "grade": 0.0,
                    "comment": ""
                }
            }
        }

    @staticmethod
    def _get_writing_template():
        return {
            "comment": "comment about student's response quality",
            "overall": 0.0,
            "task_response": {
                "Task Achievement": {
                    "grade": 0.0,
                    "comment": "comment about Task Achievement of the student's response"
                },
                "Coherence and Cohesion": {
                    "grade": 0.0,
                    "comment": "comment about Coherence and Cohesion of the student's response"
                },
                "Lexical Resource": {
                    "grade": 0.0,
                    "comment": "comment about Lexical Resource of the student's response"
                },
                "Grammatical Range and Accuracy": {
                    "grade": 0.0,
                    "comment": "comment about Grammatical Range and Accuracy of the student's response"
                }
            }
        }