"""Example: run a deep-research task with the AIR SDK.

Usage::

    export AIR_API_KEY="air_k1_..."
    python examples/deep_research.py

By default this connects to http://localhost:8000.  Override with::

    export AIR_BASE_URL="https://your-backend.example.com"
"""

import os
import sys
from dotenv import load_dotenv
import air

load_dotenv()

API_KEY = os.environ.get("AIR_API_KEY")
BASE_URL = os.environ.get("AIR_BASE_URL", "http://localhost:8000")
WORK_DIR = os.environ.get("AIR_WORK_DIR")
print(f"BASE_URL: {BASE_URL}")
print(f"WORK_DIR: {WORK_DIR}")

if not API_KEY:
    sys.exit("Set AIR_API_KEY first:  export AIR_API_KEY='air_k1_...'")

if WORK_DIR is None:
    WORK_DIR = os.getcwd()

client = air.AIR(api_key=API_KEY, base_url=BASE_URL)


task = r"""Here is a portfolio of synthetic candidate compounds across programs.
    The data is in the file /Users/boris/Desktop/data/synthetic_drug_dev_portfolio.csv
    Provide a dataguide for further analysis."""

result = client.deep_research(
    # "Write a Python script that generates noisy log-normal data and fits a log-normal distribution to the data",
    task,

    # Planning configuration
    max_plan_steps=3,
    n_plan_reviews=1,
    plan_instructions="Do two engineer steps and researcher at the end.",
    # Control configuration
    max_rounds=100,
    max_attempts=3,
    # Execution
    on_output=lambda s: print(s, end="\n", flush=True),
    work_dir="/Users/boris/Desktop/air_workdir",
    venv_path="/Users/boris/Desktop/airenv",
)

print(f"\nTask:      {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()
