import google.generativeai as genai
from dotenv import load_dotenv
import os

load_dotenv()
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

def generate_code(user_input, language):
    prompt = f"""
    Generate code in {language} that does the following:

    {user_input}

    Provide only the code, no explanations or comments unless absolutely necessary.
    Don't include any delimiters at the beginning and the ending of the code.
    """
    model=genai.GenerativeModel("gemini-1.5-flash")
    response = model.generate_content(prompt)

    return response.text

def main():
    user_input = input("Enter your query: ")

    print("Available languages (choose one):")
    print("1. Python")
    print("2. JavaScript")
    print("3. Java")
    print("4. SQL")
    print("5. C++")
    print("6. R")


    while True:
        language_choice = input("Enter the number corresponding to your preferred language: ")
        try:
            language_choice = int(language_choice)
            if 1 <= language_choice <= 6:  # Updated range
                break
            else:
                print("Invalid choice. Please enter a number from the list.")
        except ValueError:
            print("Invalid input. Please enter a number.")

    languages = ["Python", "JavaScript", "Java", "SQL", "C++", "R"]  # Updated list
    chosen_language = languages[language_choice - 1]

    generated_code = generate_code(user_input, chosen_language)
    print("\nGenerated Code:\n")
    print(generated_code)

    # (Optional) Save to file (be careful with file names from user input)
    try:
        filename = input("Enter a filename to save (or press Enter to skip): ")
        if filename:
            extension = ""
            if chosen_language == "Python":
                extension = ".py"
            elif chosen_language == "JavaScript":
                extension = ".js"
            elif chosen_language == "Java":
                extension = ".java"
            elif chosen_language == "SQL":
                extension = ".sql"
            elif chosen_language == "C++":
                extension = ".cpp"
            elif chosen_language == "R":
                extension = ".r"

            if not filename.endswith(extension):  # Simplified extension adding
                filename += extension

            with open(filename, "w") as f:
                f.write(generated_code)
            print(f"Code saved to {filename}")

    except Exception as e:
        print(f"Error saving file: {e}")

if __name__ == "__main__":
    main()