"""Test restart_at_step for project results.

Assumes a project already has idea.md and methods.md (run pipeline
examples first). The backend finds the most recent previous task and
copies its artifacts when restarting.

restart_at_step values:
    -1  = full run (planning + all control steps)
     1  = skip planning, re-run all control steps from step 1
     N  = skip planning + steps 1..N-1, re-run from step N onwards

Usage:
    # Full run:
    python 10b_results_restart.py

    # Restart from step 2:
    python 10b_results_restart.py --restart 2

    # Use a specific project:
    python 10b_results_restart.py --project my-project-name
"""

import os
import sys
import air

restart_at_step = -1
project_name = None

args = sys.argv[1:]
i = 0
while i < len(args):
    if args[i] == "--restart" and i + 1 < len(args):
        restart_at_step = int(args[i + 1])
        i += 2
    elif args[i] == "--project" and i + 1 < len(args):
        project_name = args[i + 1]
        i += 2
    else:
        i += 1

client = air.AIR(
    api_key=os.environ["AIR_API_KEY"],
    base_url=os.environ.get("AIR_BASE_URL", "http://localhost:8000"),
)

# Find a project to use
if project_name is None:
    projects = client.list_projects()
    if not projects:
        print("No projects found. Run an idea+methods example first.")
        client.close()
        sys.exit(1)
    project_name = projects[-1]

print(f"Project: {project_name}")
print(f"restart_at_step: {restart_at_step}")

project = client.get_project(project_name)

output_dir = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    f"output_results_{project_name}",
)

result = project.results(
    max_steps=2,
    max_attempts=5,
    n_plan_reviews=1,
    restart_at_step=restart_at_step,
    engineer_model="gemini-3.1-flash-lite-preview",
    researcher_model="gemini-3.1-flash-lite-preview",
    work_dir=output_dir,
)

print(f"\nResult:\n{result[:500] if result else '(empty)'}...")

client.close()
