from agent import CodingAgent
from utils import load_config

def main():
    config = load_config('configs/model_config.yaml')
    agent = CodingAgent("./final_model")

    while True:
        print("\nChoose an action:")
        print("1. Generate code")
        print("2. Answer coding question")
        print("3. Explain code")
        print("4. Suggest improvements")
        print("5. Exit")

        choice = input("Enter your choice (1-5): ")

        if choice == '1':
            prompt = input("Enter a prompt for code generation: ")
            code = agent.generate_code(prompt)
            print("\nGenerated Code:\n", code)
        elif choice == '2':
            question = input("Enter your coding question: ")
            answer = agent.answer_coding_question(question)
            print("\nAnswer:\n", answer)
        elif choice == '3':
            code = input("Enter the code to explain: ")
            explanation = agent.explain_code(code)
            print("\nExplanation:\n", explanation)
        elif choice == '4':
            code = input("Enter the code for improvement suggestions: ")
            suggestions = agent.suggest_improvements(code)
            print("\nSuggestions:\n", suggestions)
        elif choice == '5':
            break
        else:
            print("Invalid choice. Please try again.")

if __name__ == "__main__":
    main()