"""
Gemini Agents API — Exploration Script
=======================================
Purpose: Explore Google's Antigravity agent and managed agents API before
         wiring into P2HS. Run each section independently by commenting
         out what you don't need.

Prerequisites:
    pip install google-genai

Set env var before running:
    $env:GEMINI_API_KEY = "your-key-here"

Docs: https://ai.google.dev/gemini-api/docs/custom-agents
"""

import os
import json
from google import genai

# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------

API_KEY = os.environ.get("GEMINI_API_KEY")
if not API_KEY:
    raise EnvironmentError("GEMINI_API_KEY is not set.")

client = genai.Client(api_key=API_KEY)

BASE_AGENT = "antigravity-preview-05-2026"


# ---------------------------------------------------------------------------
# SECTION 1 — Ping: inline interaction with the raw Antigravity agent
# No persistence, no managed agent. Just confirms the API key works and
# the base agent responds. Antigravity requires environment="remote" (or a
# full environment dict with type/sources) on every interactions.create call.
# ---------------------------------------------------------------------------

def section_1_ping():
    print("\n=== SECTION 1: Ping Antigravity ===")

    interaction = client.interactions.create(
        agent=BASE_AGENT,
        input="Say hello and list your available default tools in one short paragraph.",
        environment="remote",
    )

    print("Output:", interaction.output_text)
    print("Environment ID:", getattr(interaction, "environment_id", None))


# ---------------------------------------------------------------------------
# SECTION 2 — Inline customization
# Overrides system instruction and injects an AGENTS.md inline.
# Good for understanding how file-based config works before committing to
# a managed agent.
# ---------------------------------------------------------------------------

def section_2_inline_customization():
    print("\n=== SECTION 2: Inline Customization ===")

    interaction = client.interactions.create(
        agent=BASE_AGENT,
        input=(
            "Create a minimal Flutter app scaffold: a single screen with a "
            "centered 'Hello P2HS' text widget. Output the content."
        ),
        system_instruction=(
            "You are a mobile app scaffolding agent. "
            "When asked to create an app, output clean, production-ready code. "
            "Always include dependencies alongside any code files."
        ),
        environment={
            "type": "remote",
            "sources": [
                {
                    "type": "inline",
                    "target": ".agents/AGENTS.md",
                    "content": (
                        "# Mobile App Agent\n\n"
                        "- Always scaffold Flutter apps unless the user specifies React Native.\n"
                        "- Output files as clearly labelled code blocks.\n"
                        "- Include a minimal pubspec.yaml with flutter SDK constraint.\n"
                        "- Never include platform-specific (android/ios) boilerplate unless asked.\n"
                    ),
                },
            ],
        },
    )

    print("Output:\n", interaction.output_text)


# ---------------------------------------------------------------------------
# SECTION 3 — Create a managed (persistent) mobile-app-builder agent
# Registers the agent by ID so it can be invoked repeatedly.
# Only need to run this once; after creation, use section_4 to invoke it.
# ---------------------------------------------------------------------------

MOBILE_AGENT_ID = "p2hs-mobile-app-builder"

def section_3_create_managed_agent():
    print("\n=== SECTION 3: Create Managed Agent ===")

    agent = client.agents.create(
        id=MOBILE_AGENT_ID,
        base_agent=BASE_AGENT,
        system_instruction=(
            "You are P2HS's mobile app builder agent. "
            "Your job is to scaffold complete, runnable Flutter applications "
            "from a user's plain-English description. "
            "Rules:\n"
            "- Default to Flutter/Dart unless the user explicitly requests React Native.\n"
            "- Always output: main.dart, pubspec.yaml, and a brief README.md.\n"
            "- Use Material 3 design and null-safe Dart.\n"
            "- Keep dependencies minimal — no unnecessary packages.\n"
            "- Structure code for readability: one widget per file when the app has >2 screens.\n"
            "- Never generate platform runner files (android/, ios/, web/) unless asked.\n"
            "- After outputting code, summarize what was created in a JSON block:\n"
            "  {\"files\": [...], \"entry_point\": \"...\", \"flutter_version\": \"...\"}\n"
        ),
        base_environment={
            "type": "remote",
            "sources": [
                {
                    "type": "inline",
                    "target": ".agents/AGENTS.md",
                    "content": (
                        "# P2HS Mobile App Builder\n\n"
                        "## Output contract\n"
                        "Always produce:\n"
                        "1. `main.dart` — entry point\n"
                        "2. `pubspec.yaml` — dependencies\n"
                        "3. `README.md` — what the app does and how to run it\n\n"
                        "## Style\n"
                        "- Material 3, null-safe Dart\n"
                        "- Minimal packages\n"
                        "- One widget per file for multi-screen apps\n\n"
                        "## Final JSON summary\n"
                        "End every response with a ```json block listing files created.\n"
                    ),
                },
            ],
        },
    )

    print(f"Created managed agent: {agent.id}")
    print(json.dumps(agent.__dict__, default=str, indent=2))


# ---------------------------------------------------------------------------
# SECTION 4 — Invoke the managed agent with a real app prompt
# This is the core test: send a natural-language app description and get
# scaffolded Flutter code back.
# ---------------------------------------------------------------------------

def section_4_invoke_managed_agent():
    print("\n=== SECTION 4: Invoke Managed Agent ===")

    result = client.interactions.create(
        agent=MOBILE_AGENT_ID,
        input=(
            "Build me a Flutter todo app with the following features:\n"
            "- A home screen listing all todos\n"
            "- Ability to add a new todo via a FAB button and a bottom sheet\n"
            "- Mark todos as complete with a checkbox\n"
            "- Delete a todo by swiping left\n"
            "- Store todos in memory (no database needed for now)\n"
            "- After creating the app, run: `flutter build web --release` inside the project directory and tell me if Flutter is available in this environment.\n"
            "Use Material 3 and keep it clean."
        ),
        environment="remote",
    )

    print("Output:\n", result.output_text)
    print("\nEnvironment ID (save for forking):", getattr(result, "environment_id", None))


# ---------------------------------------------------------------------------
# SECTION 5 — Fork the environment from section 4 into a new agent variant
# After running section 4, paste the environment_id below and run this to
# create a forked agent that starts from the already-set-up environment.
# Useful for iterating without re-running the full setup each time.
# ---------------------------------------------------------------------------

FORKED_ENV_ID = ""  # paste environment_id from section_4 output here

def section_5_fork_environment():
    if not FORKED_ENV_ID:
        print("\n=== SECTION 5: Skipped (no FORKED_ENV_ID set) ===")
        return

    print("\n=== SECTION 5: Fork Environment into New Agent ===")

    agent = client.agents.create(
        id=f"{MOBILE_AGENT_ID}-v2",
        base_agent=BASE_AGENT,
        system_instruction=(
            "You are P2HS's mobile app builder v2. "
            "Same rules as v1 but also support React Native/Expo output "
            "when the user says 'React Native'."
        ),
        base_environment=FORKED_ENV_ID,
    )

    print(f"Forked agent: {agent.id}")


# ---------------------------------------------------------------------------
# SECTION 6 — List and inspect all managed agents on this API key
# ---------------------------------------------------------------------------

def section_6_list_agents():
    print("\n=== SECTION 6: List All Managed Agents ===")

    agents = client.agents.list()
    agent_list = getattr(agents, "agents", [])

    if not agent_list:
        print("No managed agents found.")
        return

    for a in agent_list:
        print(f"  - {a.id}")

    print(f"\nTotal: {len(agent_list)} agent(s)")


# ---------------------------------------------------------------------------
# SECTION 7 — Clean up: delete the managed agent (run when done exploring)
# ---------------------------------------------------------------------------

def section_7_delete_agent():
    print("\n=== SECTION 7: Delete Managed Agent ===")
    client.agents.delete(id=MOBILE_AGENT_ID)
    print(f"Deleted: {MOBILE_AGENT_ID}")


# ---------------------------------------------------------------------------
# Runner — comment out sections you don't want to run
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Run in order. Comment out sections you've already completed.

    #section_1_ping()
    #section_2_inline_customization()
    #section_3_create_managed_agent()
    section_4_invoke_managed_agent()
    # section_5_fork_environment()   # requires FORKED_ENV_ID from section_4
    # section_6_list_agents()
    # section_7_delete_agent()       # run last when done