import guidance
from tree_of_thoughts.abstractLanguageModel import AbstractLanguageModel
import time
import os
import openai


class GuidanceLanguageModel(AbstractLanguageModel):
    def __init__(self, model, strategy="cot", evaluation_strategy="value", enable_ReAct_prompting=False):
        # gpt4 = guidance.llms.OpenAI("gpt-4")
        # vicuna = guidance.llms.transformers.Vicuna("your_path/vicuna_13B", device_map="auto")
        self.model = model
        
        # reference : https://www.promptingguide.ai/techniques/react
        self.ReAct_prompt = ''
        if enable_ReAct_prompting:
            self.ReAct_prompt = '''{{#assistant~}}
            {{gen 'Observation' temperature=0.5 max_tokens=50}}
            {{~/assistant}}'''
        
        self.strategy = strategy
        self.evaluation_strategy = evaluation_strategy
        
        self.thoughts_program = guidance('''
            {{#system~}}
            You are a logical and rational assistant.
            {{~/system}}

            {{#user~}}
            Given the current state of reasoning:
            {{state_text}}
            Generate {{k}} coherent thoughts as short as possible to continue the reasoning process.
            Don't answer the question yet.
            {{~/user}}

            %s
            
            {{#assistant~}}
            {{gen 'Thoughts' temperature=0.5 max_tokens=50}}
            {{~/assistant}}
            ''' % self.ReAct_prompt, llm=self.model)
        
        self.value_program = guidance('''
            {{#system~}}
            You are a logical and rational assistant.
            {{~/system}}

            {{#user~}}
            Given the current state of reasoning:
            {{state_text}}
            Evaluate its value as a float between 0 and 1, and NOTHING ELSE
            Don't answer the question yet.
            {{~/user}}

            {{#assistant~}}
            {{gen 'Value' temperature=1 max_tokens=10}}
            {{~/assistant}}
            ''', llm=self.model)
        
        self.vote_program = guidance('''
            {{#system~}}
            You are a logical and rational assistant.
            {{~/system}}

            {{#user~}}
            Given the following states of reasoning, vote for the best state:
            {{states_text}}
            Give the index of your voted best state(the 1st state has index 0), and NOTHING ELSE
            Don't answer the question yet.
            {{~/user}}

            {{#assistant~}}
            {{gen 'Vote' temperature=1 max_tokens=10}}
            {{~/assistant}}
            ''', llm=self.model)
        
    def model_response_handler(self, program, **kargs):
        print("Calling guidance model(Modify Me to handle specific LLM response excpetions!)")
        reponse = program(**kargs)
        return reponse

    def generate_thoughts(self, state, k):
        #implement the thought generation logic using self.model
        state_text = ' '.join(state)
        
        thoughts = []
        for _ in range(k):
            response = self.model_response_handler(self.thoughts_program, state_text=state_text, k=1)
            text = response['Thoughts']
            thoughts += [text]
        # print(thoughts)
        print(f"Generated thoughts: {thoughts}")
        return thoughts

    def evaluate_states(self, states):
        #implement state evaluation logic using self.model
        if self.evaluation_strategy == 'value':
            state_values = {}
            for state in states:
                state_text = ' '.join(state)
                response = self.model_response_handler(self.value_program, state_text=state_text)
                try:
                    value_text = response['Value']
                    print(f"Value text {value_text}")
                    value = float(value_text)
                    print(f"value: {value}")
                except ValueError:
                    value = 0  # Assign a default value if the conversion fails
                state_values[state] = value
            return state_values

        elif self.evaluation_strategy == 'vote':
            states_text = '\n'.join([' '.join(state) for state in states])
            response = self.model_response_handler(self.vote_program, states_text=states_text)
            best_state_text = response['Vote']
            print(f"Best state text: {best_state_text}")
            best_state = int(best_state_text)
            return {state: 1 if i == best_state else 0 for i in range(len(states))}

        else:
            raise ValueError("Invalid evaluation strategy. Choose 'value' or 'vote'.")
        


class GuidanceOpenAILanguageModel(GuidanceLanguageModel):
    def __init__(self, api_key, strategy="cot", evaluation_strategy="value", api_base="", api_model="", enable_ReAct_prompting=False):
        if api_key == "" or api_key == None:
            api_key = os.environ.get("OPENAI_API_KEY", "")
        if api_key != "":
            openai.api_key = api_key
        else:
            raise Exception("Please provide OpenAI API key")

        if api_base == ""or api_base == None:
            api_base = os.environ.get("OPENAI_API_BASE", "")  # if not set, use the default base path of "https://api.openai.com/v1"
        if api_base != "":
            # e.g. https://api.openai.com/v1/ or your custom url
            openai.api_base = api_base
            print(f'Using custom api_base {api_base}')
            
        if api_model == "" or api_model == None:
            api_model = os.environ.get("OPENAI_API_MODEL", "")
        if api_model != "":
            self.api_model = api_model
        else:
            self.api_model = "text-davinci-003"
        print(f'Using api_model {self.api_model}')

        super().__init__(guidance.llms.OpenAI(self.api_model), strategy, evaluation_strategy, enable_ReAct_prompting)
        
    
    def model_response_handler(self, program, **kargs):
        error_msg = ''
        while True:
            try:
                program.llm.max_retries = 60
                guidance.llms.OpenAI.cache.clear()
                response = program(**kargs)
                return response
            except openai.error.RateLimitError as e:
                sleep_duratoin = os.environ.get("OPENAI_RATE_TIMEOUT", 30)
                print(f'{str(e)}, sleep for {sleep_duratoin}s, set it by env OPENAI_RATE_TIMEOUT')
                time.sleep(sleep_duratoin)
            except Exception as e:
                if str(e) == f'''Too many (more than {guidance.llm.max_retries}) OpenAI API RateLimitError's in a row!''':
                    sleep_duratoin = os.environ.get("OPENAI_RATE_TIMEOUT", 30)
                    print(f'{str(e)}, sleep for {sleep_duratoin}s, set it by env OPENAI_RATE_TIMEOUT')
                    time.sleep(sleep_duratoin)
                else:
                    error_msg = str(e)
                    break
        raise Exception(error_msg)
