#!/usr/bin/env python3
"""
End-to-end pipeline: idea -> literature -> methods -> results -> paper -> review.

The results step uses a WebSocket session with local code execution
(plots are generated on your machine, then uploaded to the server).

Usage:
    export AIR_API_KEY="air_k1_..."
    python examples/end_to_end.py

    # or point at a remote server:
    AIR_BASE_URL="http://localhost:8000" python examples/end_to_end.py

    or run the air-backend locally and point to AIR_BASE_URL="http://127.0.0.1:8000"
"""

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")

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

PROJECT_NAME = "end2end_project"
print(f"Running end-to-end test with project name: {PROJECT_NAME}")
print(f"BASE_URL: {BASE_URL}")

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

    project = client.create_project(
        name=PROJECT_NAME,
        data_description="We want to do a basic study of simple harmonic oscillator.",
    )
    print(f"Project created: {project}")

    print("\n[1/6] Generating research idea ...")
    idea = project.idea()
    print(f"  Idea: {idea}")

    print("\n[2/6] Running literature search ...")
    lit = project.literature()
    print(f"  Literature: {lit}")

    print("\n[3/6] Developing methods ...")
    meth = project.methods()
    print(f"  Methods: {meth}")

    print("\n[4/6] Generating results (local code execution via WebSocket) ...")
    res = project.results(
        max_attempts=10,
        max_steps=2,
        on_output=lambda s: print(s, end="\n", flush=True),
    )
    print(f"  Results: {len(res)} chars")

    print("\n[5/6] Writing paper ...")
    tex = project.paper(journal=None, add_citations=True)
    print(f"  Paper: {len(tex)} chars")

    print("\n[6/6] Running review ...")
    rev = project.review()
    print(f"  Review: {len(rev)} chars")

    # Pull all files produced to default to "~/ai-scientist" or set AIR_LOCAL_DIR) to local disk
    project.pull_files()
