import os
import subprocess
from datetime import datetime

# --- Configuration ---
MODEL = "llama3"
# MODEL = "deepseek-coder"
OUTPUT_DIR = "generated_code"
TARGETS = ["Python3"]

PROMPT_TEMPLATE = """
Generate {count} different working code examples in {framework}.
Task: {task}.
Each example should be complete and runnable with Python 3.
"""

def generate_code(task, count=3):
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    for framework in TARGETS:
        prompt = PROMPT_TEMPLATE.format(
            count=count, framework=framework, task=task
        )
        print(f"[+] Requesting {framework} code...")

        result = subprocess.run(
            ["ollama", "run", MODEL],
            input=prompt.encode(),
            capture_output=True,
        )
        output = result.stdout.decode().strip()

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{OUTPUT_DIR}/{framework}_{timestamp}.txt"
        with open(filename, "w") as f:
            f.write(output)
        print(f"[✓] Saved: {filename}")

def run_tasks(taskfile, count):
    with open(taskfile, "r") as f:
        tasks = [line.strip() for line in f if line.strip()]

    for task in tasks:
        print(f"\n=== Generating code for: {task} ===")
        generate_code(task=task, count=count)

if __name__ == "__main__":
    run_tasks("python_tasks.txt", count=5)