#!/usr/bin/env python3
"""
Generate an idea, then demonstrate two independent ways to push edits back.

  Edit A — in-memory edit pushed with project.write_file()
  Edit B — local file edit pushed with project.push_files()

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

import os
import textwrap
from datetime import datetime, timezone
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")

IDEA_PATH = "Iteration0/input_files/idea.md"
LOCAL_DIR = "idea_edit_example"

DATA_DESCRIPTION = """
    Star catalog dataset (Hipparcos) with astrometric and photometric features
    for 100 000+ stars: positions, parallaxes, proper motions, magnitudes, colours.
    File: hipparcos-voidmain.csv
"""


def apply_edit_a(original: str) -> str:
    """Edit A: prepend a human-review banner (scope narrowed to kinematics)."""
    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    banner = textwrap.dedent(f"""\
        > **[Edit A — write_file]** Reviewed on {timestamp}.
        > Scope narrowed to kinematic clustering (parallax + proper motion);
        > photometric features excluded.

    """)
    return banner + original


def apply_edit_b(text: str) -> str:
    """Edit B: append a follow-up note about evaluation strategy."""
    note = textwrap.dedent("""\

        ---
        > **[Edit B — push_files]** Follow-up annotation added from local file.
        > Evaluation should use silhouette score on the reduced proper-motion space,
        > not raw Euclidean distance on the full feature set.
    """)
    return text + note


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

    # 1. Create project and generate idea
    project = client.create_project(
        name="idea-edit-demo",
        data_description=DATA_DESCRIPTION,
    )
    print(f"Created: {project}")

    print("\n[1/5] Generating idea …")
    idea_original = project.idea(mode="fast", timeout=900)
    print(f"  Original idea: {len(idea_original)} chars")
    print(idea_original[:300] + " …")

    # 2. Confirm server content matches the returned value
    print("\n[2/5] Reading idea back from server to confirm …")
    server_idea = project.get_file(IDEA_PATH)
    assert server_idea == idea_original, "Server content differs from returned value!"
    print(f"  OK — {len(server_idea)} chars on server")

    # ------------------------------------------------------------------ #
    # Edit A: manipulate the string in memory, push with write_file()     #
    # ------------------------------------------------------------------ #
    print("\n[3/5] Edit A — prepend review banner, push with write_file() …")
    idea_after_a = apply_edit_a(server_idea)
    project.write_file(IDEA_PATH, idea_after_a)
    print(f"  Pushed {len(idea_after_a)} chars (added {len(idea_after_a) - len(server_idea)} chars)")

    # Verify Edit A landed on the server
    server_after_a = project.get_file(IDEA_PATH)
    assert "[Edit A — write_file]" in server_after_a, "Edit A not found on server!"
    print("  Verified — Edit A is present on server")

    # ------------------------------------------------------------------ #
    # Edit B: pull to disk, edit the local file, push with push_files()  #
    # ------------------------------------------------------------------ #
    print("\n[4/5] Edit B — pull locally, append evaluation note, push with push_files() …")
    written = project.pull_files(local_dir=LOCAL_DIR)
    local_idea_path = next(p for p in written if p.endswith("idea.md"))
    print(f"  Pulled to: {local_idea_path}")

    with open(local_idea_path, encoding="utf-8") as fh:
        local_content = fh.read()

    idea_after_b = apply_edit_b(local_content)
    with open(local_idea_path, "w", encoding="utf-8") as fh:
        fh.write(idea_after_b)
    print(f"  Local file updated: {len(local_content)} → {len(idea_after_b)} chars")

    uploaded = project.push_files(local_idea_path, remote_path=IDEA_PATH)
    print(f"  Pushed {len(uploaded)} chars via push_files()")

    # Verify Edit B landed on the server
    server_after_b = project.get_file(IDEA_PATH)
    assert "[Edit A — write_file]" in server_after_b, "Edit A missing after Edit B!"
    assert "[Edit B — push_files]" in server_after_b, "Edit B not found on server!"
    print("  Verified — both edits are present on server")

    # 5. Show final state
    print("\n[5/5] Final idea on server (first 600 chars):")
    print("-" * 60)
    print(server_after_b[:600])
    print("…")
    print("-" * 60)
    print(f"\nDone. Both edit methods confirmed working.")
    print(f"  Edit A (write_file):  in-memory string → server")
    print(f"  Edit B (push_files):  local file → server")
