import re
from typing import List, Dict
from config import *
from prompt import *
from llm import OpenAI3, OpenAI4
import random


def delete_index(string: str) -> str:
    if len(string) > 1:  # incase of the sting is a single letter when generate word complete question
        if string[1] in ['.', ':', '、', '：', ')', '）']:
            string = string[2:]
    return string


def normalize_options_and_answer(answer, options):
    if len(options) != 1:
        random.shuffle(options)
        options = [delete_index(option) for option in options]
        answer = delete_index(answer)
        answer_index = options.index(answer)
        answer = number_letter_dict[answer_index] + ':' + answer
        options = [number_letter_dict[i] + ':' + options[i] + ',' for i in range(len(options))]
        options = ''.join(options)
        options = options[:-1]
    else:
        options = options[0]

    return answer, options


def response_to_question(response):
    try:
        try:
            _, stem, options, answer = re.split(
                r'stem: |stem:|Stem: |Stem:'
                r'|options: |options:|Options: |Options:'
                r'|Correct answer: |Correct answer:|correct answer: |correct answer:'
                r'|answer: |answer:|Answer: |Answer:'
                , response)
        except:

            try:
                _, stem, options = re.split(
                    r'stem: |stem:|Stem: |Stem:'
                    r'|options: |options:|Options: |Options:'
                    , response)
                answer = None

            except:
                _, question, answer = re.split(
                    r'question: |question:|Question: |Question:'
                    r'|Correct answer: |Correct answer:|correct answer: |correct answer:'
                    r'|answer: |answer:|Answer: |Answer:'
                    , response)
                stem = None
                options = None
                while question[-1] == '\n':
                    question = question[:-1]

        if stem:
            while stem[-1] == '\n':
                stem = stem[:-1]

        if options:
            options = options.split('\n')
            options = [option for option in options if option != '']

            if answer is None:
                answer = 'None'
                for i in range(len(options)):
                    if 'correct answer' in options[i]:
                        answer = options[i]
                        answer = re.split(r'\(correct answer\)|correct answer', answer)
                        answer = [i for i in answer if i is not None][0]
                        options[i] = answer

            answer = answer.split('\n')[0]
            answer, options = normalize_options_and_answer(answer, options)
            question = stem + '\n' + options

    except Exception as e:
        print(f'fail!,reason:{e},response:{response}')
        question = 'fail!'
        answer = 'fail!'

    return question, answer


def response_to_options(response):

    options = response.split('\n')
    options = [delete_index(option) for option in options]

    return options


def generate_word_complete_question(words: List[str], options_number: int = 4) -> List[Dict]:
    responses = []
    for word in words:

        vowels = []
        for vowel in vowel_Combination:
            if vowel in word:
                vowels.append(vowel)

        if vowels:
            answer = vowels[random.randint(0, len(vowels)-1)]

            stem = ''
            _stem = word.split(answer)
            index = random.randint(0, len(_stem)-2)
            for i in range(len(_stem)-1):
                if i != index:
                    stem = stem + _stem[i] + answer
                else:
                    stem = stem + _stem[i] + '_' * len(answer)
            stem = stem + _stem[-1]

            options = [i for i in vowel_Combination if (len(i) == len(answer)) & (i != answer)]
            if len(options) > options_number-1:
                random.shuffle(options)
                options = options[:options_number-1]
            options.append(answer)

        else:
            index = list(range(len(word)))
            num = random.randint(2, 3)
            num = len(word)-3 if num > len(word)-3 else num

            random.shuffle(index)
            ans_index = index[:num]
            ans_index.sort()
            answer = [word[i]+',' for i in ans_index]
            answer = ''.join(answer)
            answer = answer[:-1]

            stem = ['_' if i in ans_index else word[i] for i in range(len(word))]
            stem = ''.join(stem)

            options = [answer]
            count = 1
            while count < options_number:
                random.shuffle(index)
                option_index = index[:num]
                option = [word[i] + ',' for i in option_index]
                option = ''.join(option)
                option = option[:-1]
                if option not in options:
                    options.append(option)
                    count += 1

        answer, options = normalize_options_and_answer(answer, options)

        question = stem + '\n' + options

        responses.append({'question': question, 'answer': answer})

    return responses


def generate_sentence_recomposition_question(sentences: List[str], generate_mode: bool =False) -> List[Dict]:
    responses = []
    for sentence in sentences:
        try:
            prompt = 'f"' + similar_sentence_generate_prompt + '"'
            if generate_mode:
                try:
                    prompt = eval(prompt)
                    sent = _generate(prompt)
                    sent = re.split(r'[.?!]', sent)[0]

                except:
                    sent = sentence
            else:
                sent = sentence
            answer = sent
            words = sent.split(' ')
            stem = '______ ' * len(words)
            stem = stem[:-1]
            symbol = words[-1][-1]
            if symbol in ['.', '?', '!', '。', '？', '！']:
                stem = stem + symbol
                words[-1] = words[-1][:-1]

            words_copy = words.copy()
            while words_copy == words:
                random.shuffle(words)
            words = ','.join(words)

            question = stem + '\n' + words
            responses.append({'question': question, 'answer': answer})
        except Exception as e:
            print(e)

    return responses


def _generate(prompt, model='openai_3'):
    model = eval(model_name[model])()
    response = model(prompt)
    return response


def generate_phrase_select_question(phrases: List[str], options_number: int = 2) -> List[Dict]:
    responses = []
    num = options_number - 1
    for phrase in phrases:
        try:
            prompt = 'f"' + similar_phrase_generate_prompt + '"'
            phrase = phrase[0].upper() + phrase[1:]
            prompt = eval(prompt)
            response = _generate(prompt)
            options = response_to_options(response)
            options.append(phrase)
            answer, options = normalize_options_and_answer(phrase, options)
            # responses.append({'stem': None, 'options': options, 'answer': answer})
            responses.append({'question': options, 'answer': answer})
        except Exception as e:
            print(e)
    return responses


def generate_sentence_complete_question(sentences: List[str], options_number: int = 4) -> List[Dict]:
    responses = []
    num = options_number - 1
    for sentence in sentences:
        try:
            prompt = 'f"' + sentence_complete_question_generate_prompt + '"'
            prompt = eval(prompt)
            prompt = prompt.replace('.', '\n')

            response = _generate(prompt)
            question, answer = response_to_question(response)
            responses.append({'question': question, 'answer': answer})
        except Exception as e:
            print(e)
    return responses


def generate_grammar_question(grammars: List[str],  options_number: int = 4) -> List[Dict]:
    responses = []
    num = options_number - 1
    for grammar in grammars:
        try:
            prompt = 'f"' + grammar_question_generate_prompt + '"'
            prompt = eval(prompt)
            prompt = prompt.replace('.', '\n')
            response = _generate(prompt)
            question, answer = response_to_question(response)
            responses.append({'question': question,  'answer': answer})
        except Exception as e:
            print(e)
    return responses

ans1 = generate_word_complete_question(['word', 'manager', 'answer', 'fight', 'jump', 'hihhttgrh'])

# ans2 = generate_phrase_select_question(['the Dragon Boat Festival'], 3)

# ans3 = generate_sentence_recomposition_question(['You really require a lot of talent and hard work to succeed.',
# 'I have a job interview tomorrow.], False)

# ans4 = generate_sentence_complete_question(['You really require a lot of talent and hard work to succeed.',
#                                             'I have a job interview tomorrow.'])

# ans5 = generate_grammar_question(['be supposed to', 'It is + adj. + 动词不定式'])
# print(ans5)