"""Test restart_at_step for deep_research.

The backend automatically finds the most recent previous task for the
user and copies its artifacts (plan, context pickles, control outputs)
into the new task directory. The client always creates a fresh task_id.

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

Usage:
    # Full run (planning + control):
    python 09b_deep_research_restart.py

    # Skip planning, re-run control from step 1:
    python 09b_deep_research_restart.py --restart 1

    # Skip planning + step 1, re-run from step 2:
    python 09b_deep_research_restart.py --restart 2
"""

import os
import sys
import air

restart_at_step = -1
if "--restart" in sys.argv:
    idx = sys.argv.index("--restart")
    if idx + 1 < len(sys.argv):
        restart_at_step = int(sys.argv[idx + 1])
        print(f"Restarting from step {restart_at_step}")

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

task = """
1. Write a Python script that prints "Hello from step 1" and saves it to step1_output.txt
2. Write a Python script that reads step1_output.txt and prints its contents with "Step 2 done" appended
"""

output_dir = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "output_deep_research_restart",
)

print(f"Task:\n{task.strip()}\n")
print(f"restart_at_step: {restart_at_step}")
print(f"work_dir: {output_dir}\n")

result = client.deep_research(
    task=task,
    max_plan_steps=2,
    n_plan_reviews=1,
    restart_at_step=restart_at_step,
    plan_instructions="1. engineer: write and run step 1, 2. engineer: write and run step 2.",
    engineer_model="gemini-3.1-flash-lite-preview",
    researcher_model="gemini-3.1-flash-lite-preview",
    max_rounds=50,
    max_attempts=3,
    work_dir=output_dir,
)

print(f"\nTask ID:  {result.task_id}")
print(f"Work dir: {result.work_dir}")
print(f"Files:    {[f.path for f in result.files_created]}")
if result.error:
    print(f"Error:    {result.error}")

client.close()
