import openai 
import re
import subprocess

# openai.api_key = "YOUR_API_KEY"
# OPENAI_API_KEY in your env varaBles


response = openai.ChatCompletion.create(
  model="gpt-3.5-turbo",
  stream=True,
  messages=[
        {"role": "system", "content": "You are a expert python programmer. ALWAYS RETURN THE CODE STARTING WITH ```python."},
        {"role": "user", "content": "write a python function which calculates nth fibonacci number."},
    ]
)

# print(response['choices'][0]['message']['content'])

responses = ""
# Process each chunk
for chunk in response:
    if "role" in chunk["choices"][0]["delta"]:
        continue
    elif "content" in chunk["choices"][0]["delta"]:
        r_text = chunk["choices"][0]["delta"]["content"]
        responses += r_text
        print(r_text, end="", flush=True)

print("TOTAL RESPONSES: ", responses)

# extract the python code from within ```python and ``` 
code = re.search(r"```python(.+?)```", responses, re.DOTALL).group(1)

# save the code to a file
with open("code.py", "w") as f:
    f.write(code)

# ALWAYS BE CAREFUL WHILE AUTO RUNNING CODE FROM GPT. 
    # Run the code and capture the output and error
try:
    completed_process = subprocess.run(["python", "code.py"], capture_output=True, text=True)
    output = completed_process.stdout
    error = completed_process.stderr
except subprocess.CalledProcessError as e:
    output = e.output
    error = e.stderr

# Print the output
print("Output:")
print(output)

# Print the error, if any
if error:
    print("Error:")
    print(error)

# # read the code from the file
# with open("code.py", "r") as f:
#     code = f.read()

response = openai.ChatCompletion.create(
  model="gpt-3.5-turbo",
  stream=True,
  messages=[
        {"role": "system", "content": """You are a expert python programmer. 
                                    You are provided with some code and its output of the program and the error if any.
                                    consider the output and error and rewrite the code to handle the error if necessary
                                    always return the code starting with ```python."""},
        {"role": "user", "content": f"""
        code:

        {code}

        output:

        {output}

        error:

        {error}
        ."""},
    ]
)

# print(response['choices'][0]['message']['content'])

responses = ""
# Process each chunk
for chunk in response:
    if "role" in chunk["choices"][0]["delta"]:
        continue
    elif "content" in chunk["choices"][0]["delta"]:
        r_text = chunk["choices"][0]["delta"]["content"]
        responses += r_text
        print(r_text, end="", flush=True)


# extract the python code from within ```python and ``` 
code_fix = re.search(r"```python(.+?)```", responses, re.DOTALL).group(1)

# save the code to a file
with open("code_fix.py", "w") as f:
    f.write(code_fix)