import openai

# Put your OpenAI API key here
openai.api_key = "YOUR_OPENAI_API_KEY"

def ask_about_code(code, question):
    prompt = f"""
You are an AI assistant that helps with Python code.

Here is the Python code:
{code}

Answer this question about the code:
{question}
"""

    response = openai.Completion.create(
        engine="text-davinci-003",  # You can use GPT-4 if available
        prompt=prompt,
        max_tokens=300,
        temperature=0,
        n=1,
        stop=None,
    )
    return response.choices[0].text.strip()

def main():
    print("Welcome to the Python code assistant chatbot!")
    print("Paste your Python code first (type 'END' on a new line when done):")

    # Collect multiline code input
    code_lines = []
    while True:
        line = input()
        if line.strip().upper() == "END":
            break
        code_lines.append(line)
    code = "\n".join(code_lines)

    print("\nNow you can ask questions about your code. Type 'exit' to quit.")
    while True:
        question = input("Your question: ")
        if question.lower() == "exit":
            print("Goodbye!")
            break
        answer = ask_about_code(code, question)
        print("AI:", answer)

if __name__ == "__main__":
    main()