import os, openai, time
from tools import printInColor
from concurrent.futures import ThreadPoolExecutor, TimeoutError

openai.api_key = os.environ["OPENAI_API_KEY"]

def ask_openai(prompt, assistant=None):
    printInColor("Asking OpenAI...\n===================\n"+ prompt + "\n====================", 'green')
    attempt_limit = 3
    wait_time = 5000
    timeout_duration = 30  # seconds
    start_time = time.time() # measure the start time

    prompt = [
        {"role": "user", "content": prompt} 
    ]
    
    prompt.append({"role": "system", "content": "我们在模拟密室逃脱游戏，你是玩家，我是游戏主持人。你需要通过观察周围的环境来交互。每次玩家行为结束，我需要从头考虑环境变化和影响，构筑新的线索，调整游戏难度，重估剩余步骤。"})

    if assistant:
        prompt.append({"role": "assistant", "content": assistant})

    models = {
        "gpt4": "gpt-4",
        "turbo": "gpt-3.5-turbo",
    }

    currentModel = models["turbo"]
    while attempt_limit > 0:
        try:
            
            with ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(
                    openai.ChatCompletion.create,
                    model = currentModel,
                    messages = prompt,
                    stop = ["玩家：", "你可以", "你需要", "你有"],
                    temperature = 0.2,
                )
                response = future.result(timeout=timeout_duration)

            print(f"Token: {response.usage}")
            duration = time.time() - start_time # calculate the duration
            printInColor(f"Time use: {duration}ms", 'yellow')
            # print cost
            # calculate_cost(response, currentModel)
            break

        except TimeoutError:
            print(f"Request timed out after {timeout_duration} seconds. Attempts {attempt_limit} left.")
        except Exception as e:
            print('Failed to send request. Exception:', e)
        
        time.sleep(wait_time/1000)
        wait_time *= 2
        attempt_limit -= 1

    
    return response