import random
from typing import List
import re

from prompts.gsm8k import Decompose_Question, Decompose_Answer, Full_Decompose_Answer, CoT_Question, CoT_Answer, System_Prompt

def SingleTurn_decompose_instuction(question: str, n_shot: int, system_choice: int = 0) -> dict:
    def question_prompt(s):
        return f'Question: {s}'
    
    answer_list = Full_Decompose_Answer if system_choice == 1 else Decompose_Answer
    prompt = ''
    for i in range(n_shot):
        prompt += question_prompt(Decompose_Question[i]) + '\n'
        prompt += answer_list[i] + '\n\n'
    prompt += question_prompt(question) + '\n' + "Let's break down this problem: "

    chats = []
    # chats.append({"role": "system", "content": System_Prompt[system_choice]})
    chats.append({"role": "system", "content": System_Prompt[0]})
    chats.append({"role": "user", "content": prompt})
    return chats

def SingleTurn_cot_instuction(question: str, n_shot: int) -> dict:
    def question_prompt(s):
        return f'Question: {s}'
    
    def answer_prompt(s):
        return f'Answer: {s}'
    
    prompt = ''
    for i in range(n_shot):
        prompt += question_prompt(CoT_Question[i]) + '\n'
        prompt += answer_prompt(CoT_Answer[i]) + '\n\n'
    prompt += question_prompt(question) + '\n' + "Let's think step by step. You MUST write the final answer only as an integer after the phrase 'So the answer is'. "

    chats = []
    chats.append({"role": "user", "content": prompt})
    return chats

def SingleTurn_decompose_prompt(question: str, intermediate_response: str, n_shot:int) -> dict:
    def question_prompt(s):
        return f'Question: {s}'
    
    chats = []
    prompt = ''
    
    for i in range(n_shot):
        prompt += question_prompt(Decompose_Question[i]) + '\n'
        prompt += Full_Decompose_Answer[i] + '\n\n'
    prompt +=  question_prompt(question) + '\n' + intermediate_response + " Let's think step by step. You MUST write the final answer only as an integer after the phrase 'So the answer is'."
    chats.append({"role": "user", "content": prompt})

    return chats