from flask import Flask, request
from flask_jwt_extended import JWTManager, jwt_required
from functools import reduce
from helper.token_counter import count_tokens
from helper.openai_interface import make_openai_call
import os

from dotenv import load_dotenv

load_dotenv()

app = Flask(__name__)

app.config['JWT_SECRET_KEY'] = os.getenv("JWT_SECRET_KEY")
jwt = JWTManager(app)

GRADING_TEMPERATURE = 0.1
GEN_QUESTION_TEMPERATURE = 0.7
WRITING_TASK_2_POST_FIELDS = ['overall', 'comment', 'task_response']
WRITING_TASK_2_GET_FIELDS = ['question']

@app.route('/writing_task2', methods=['POST'])
@jwt_required()
def grade_writing_task():
    data = request.get_json()  # Assuming the request data is in JSON format
    question = data.get('question')
    answer = data.get('answer')
    messages = [
        {
            "role": "system",
            "content": "You are a IELTS examiner.",
        },
        {
            "role": "system",
            "content": f"The question you have to grade is of type Writing Task 2 and is the following: {question}",
        },
        {
            "role": "user",
            "content": "It is mandatory for you to provide your response with the overall grade and breakdown grades, "
                       "in the following json format: {'comment': 'comment about answer quality', 'overall': 7.0, 'task_response': {'Task Achievement': 8.0, "
                       "'Coherence and Cohesion': 6.5, 'Lexical Resource': 7.5, 'Grammatical Range and Accuracy': "
                       "6.0}}",
        },
        {
            "role": "user",
            "content": "Example output: { 'comment': 'Overall, the response is good but there are some areas that need "
                       "improvement.\n\nIn terms of Task Achievement, the writer has addressed all parts of the question "
                       "and has provided a clear opinion on the topic. However, some of the points made are not fully "
                       "developed or supported with examples.\n\nIn terms of Coherence and Cohesion, there is a clear "
                       "structure to the response with an introduction, body paragraphs and conclusion. However, there "
                       "are some issues with cohesion as some sentences do not flow smoothly from one to another.\n\nIn "
                       "terms of Lexical Resource, there is a good range of vocabulary used throughout the response and "
                       "some less common words have been used effectively.\n\nIn terms of Grammatical Range and Accuracy, "
                       "there are some errors in grammar and sentence structure which affect clarity in places.\n\nOverall, "
                       "this response would score a band 6.5.', 'overall': 6.5, 'task_response': "
                       "{ 'Coherence and Cohesion': 6.5, 'Grammatical Range and Accuracy': 6.0, 'Lexical Resource': 7.0, "
                       "'Task Achievement': 7.0}}",
        },
        {
            "role": "user",
            "content": f"Evaluate this answer according to ielts grading system: {answer}",
        },
    ]
    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 = make_openai_call(messages, token_count, WRITING_TASK_2_POST_FIELDS, GRADING_TEMPERATURE)
    return response

@app.route('/writing_task2', methods=['GET'])
@jwt_required()
def get_writing_task_question():
    messages = [
        {
            "role": "system",
            "content": "You are a IELTS program that generates questions for the exams.",
        },
        {
            "role": "system",
            "content": "The question you have to generate is of type Writing Task 2 and is the following.",
        },
        {
            "role": "user",
            "content": "It is mandatory for you to provide your response with the question "
                       "in the following json format: {'question': 'question'}",
        },
        {
            "role": "user",
            "content": "Example output: { 'question': 'We are becoming increasingly dependent on computers. "
                       "They are used in businesses, hospitals, crime detection and even to fly planes. What things will "
                       "they be used for in the future? Is this dependence on computers a good thing or should we he more "
                       "auspicious of their benefits?'}",
        },
        {
            "role": "user",
            "content": "Generate a question for IELTS exam Writing Task 2.",
        },
    ]
    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 = make_openai_call(messages, token_count, WRITING_TASK_2_GET_FIELDS, GEN_QUESTION_TEMPERATURE)
    return response


if __name__ == '__main__':
    app.run()