main.py

import asyncio
import sys
from llm import call_llm
from utils import clean_markdown_code

async def main():
    ascii_art = r"""
    ___                         __       ___                    __ 
   /   |  ____ _   _____  _____/ /_     /   | ____ ____  ____  / /_
  / /| | / __ \ | / / _ \/ ___/ __ \   / /| |/ __ `/ _ \/ __ \/ __/
 / ___ |/ / / / |/ /  __(__  ) / / /  / ___ / /_/ /  __/ / / / /_  
/_/  |_/_/ /_/|___/\___/____/_/ /_/  /_/  |_\__, /\___/_/ /_/\__/  V1.0
                                           /____/                    
    """
    print("=" * 75)
    print(ascii_art)
    print("=" * 75)
    
    while True:
        try:
            # Get user input
            sys.stdout.write("\nEnter coding task (or 'exit' to quit): \n> ")
            sys.stdout.flush()
            
            task = sys.stdin.readline().strip()
            
            if not task:
                continue
                
            if task.lower() in ('exit', 'quit', 'q'):
                print("Exiting...")
                break
                
            if task.startswith("save "):
                print("Error: Provide a task first, then you can save the output.")
                continue

            print("\nGenerating code...")
            
            # Call LLM
            response = await call_llm(task)
            
            # Clean response
            code = clean_markdown_code(response)
            
            print("\n==================== CODE ====================")
            print(code)
            print("==============================================\n")
            
            # Ask if user wants to save
            sys.stdout.write("Type 'save filename.py' to save, or press Enter to continue: ")
            sys.stdout.flush()
            
            save_cmd = sys.stdin.readline().strip()
            if save_cmd.startswith("save "):
                filename = save_cmd.split(" ", 1)[1].strip()
                if filename:
                    with open(filename, "w") as f:
                        f.write(code)
                    print(f"Saved to {filename}")
                else:
                    print("Invalid filename.")
                    
        except KeyboardInterrupt:
            print("\nExiting...")
            break
        except Exception as e:
            print(f"\nAn error occurred: {e}")

if __name__ == "__main__":
    asyncio.run(main())