# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install --upgrade --quiet --force-reinstall "google-cloud-aiplatform>=1.112.0" anthropic google-adk

# ✅ Installation complete!
print("✅ Packages installed successfully!")

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

# Use the environment variable if the user doesn't provide Project ID.
import os

import vertexai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

# ADK env variables
GOOGLE_GENAI_USE_VERTEXAI = 1

os.environ["GOOGLE_CLOUD_PROJECT"] = PROJECT_ID
os.environ["GOOGLE_CLOUD_LOCATION"] = LOCATION
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = str(GOOGLE_GENAI_USE_VERTEXAI)

# Initialize Vertex AI
vertexai.init(project=PROJECT_ID, location=LOCATION)
client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

# ✅ Confirmation
print("✅ Vertex AI initialized successfully!")
print(f"   Project: {PROJECT_ID}")
print(f"   Location: {LOCATION}")

import json
import logging
import re
from io import BytesIO

import matplotlib.pyplot as plt
from anthropic import AnthropicVertex
from vertexai import types
from vertexai.generative_models import (
    Content,
    FunctionDeclaration,
    GenerationConfig,
    GenerativeModel,
    Part,
    Tool,
)

logging.getLogger().setLevel(logging.INFO)
from google.adk.agents import LlmAgent
from google.adk.artifacts import InMemoryArtifactService
from google.adk.code_executors import BuiltInCodeExecutor
from google.adk.code_executors.agent_engine_sandbox_code_executor import (
    AgentEngineSandboxCodeExecutor,
)
from google.adk.events import Event
from google.adk.memory import InMemoryMemoryService
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types as genai_types
from pydantic import BaseModel, Field

def parse_event(event: Event):
    """Parse agent events to highlight function calls and code execution."""
    if event.content and event.content.parts:
        for part in event.content.parts:
            # Check for function call (agent using tool)
            if hasattr(part, "function_call") and part.function_call:
                print(f"\nTOOL CALL: {part.function_call.name}")
                if "code" in part.function_call.args:
                    print("Code to execute:")
                    print(part.function_call.args["code"].strip())

            # Check for function response (tool result)
            elif hasattr(part, "function_response") and part.function_response:
                resp = part.function_response.response
                if isinstance(resp, dict) and resp.get("status") == "success":
                    print("\nEXECUTION RESULT:")
                    print(resp.get("output", "").strip())

            # Check for Code Interpreter Extension executable code
            elif hasattr(part, "executable_code") and part.executable_code:
                print("\nCODE INTERPRETER EXECUTION:")
                print(f"Language: {part.executable_code.language}")
                print("Code:")
                print(part.executable_code.code.strip())

            # Check for Code Interpreter Extension execution result
            elif hasattr(part, "code_execution_result") and part.code_execution_result:
                print("\nCODE INTERPRETER RESULT:")
                print(f"Status: {part.code_execution_result.outcome}")
                if part.code_execution_result.output:
                    print("Output:")
                    print(part.code_execution_result.output.strip())
                if (
                    hasattr(part.code_execution_result, "error")
                    and part.code_execution_result.error
                ):
                    print(f"Error: {part.code_execution_result.error}")

            # Check for text responses (explanatory text or final response)
            elif hasattr(part, "text") and part.text:
                # For final responses, show with special formatting
                if event.is_final_response():
                    print("\nAGENT RESPONSE:")
                    print(part.text.strip())
                # For intermediate text (explanations before code)
                elif len(part.text.strip()) > 0:
                    print("\nEXPLANATION:")
                    print(part.text.strip())

# Create an AgentEngine resource (top-level container)
agent_engine = client.agent_engines.create()

print("✅ Agent Engine created successfully!")
print(f"   Resource name: {agent_engine.api_resource.name}")

sandbox_operation = client.agent_engines.sandboxes.create(
    name=agent_engine.api_resource.name,
    config=types.CreateAgentEngineSandboxConfig(display_name="my_first_sandbox"),
    spec={"code_execution_environment": {}},
)

# Save the sandbox resource name for later use
sandbox_resource_name = sandbox_operation.response.name

print("✅ Sandbox created successfully!")
print("   Display name: my_first_sandbox")
print(f"   Resource name: {sandbox_resource_name}")
print(f"   State: {sandbox_operation.response.state}")

# Execute simple Python code in the sandbox
response = client.agent_engines.sandboxes.execute_code(
    name=sandbox_resource_name,
    input_data={
        "code": "import math\nprint(f'Square root of 15376: {math.sqrt(15376)}')"
    },
)

print("✅ Code executed successfully!")

# Parse the response to extract stdout/stderr
# The JSON output (with stdout/stderr) has mime_type="application/json" and no metadata
for output in response.outputs:
    if output.mime_type == "application/json" and output.metadata is None:
        # Decode the bytes to string
        result = json.loads(output.data.decode("utf-8"))

        # Display stdout (standard output)
        if result.get("msg_out"):
            print("📤 Output:")
            print(result.get("msg_out"))

        # Display stderr (errors) if any
        if result.get("msg_err"):
            print("❌ Errors:")
            print(result.get("msg_err"))

# Define the configuration parameters right before use
# fmt: off
language_config = "LANGUAGE_PYTHON"  # @param ["LANGUAGE_UNSPECIFIED", "LANGUAGE_PYTHON", "LANGUAGE_JAVASCRIPT"] {type:"string"}
machine_config = "MACHINE_CONFIG_VCPU4_RAM4GIB"  # @param ["MACHINE_CONFIG_UNSPECIFIED", "MACHINE_CONFIG_VCPU4_RAM4GIB"] {type:"string"}
# fmt: on

# Create the custom sandbox
sandbox_operation = client.agent_engines.sandboxes.create(
    name=agent_engine.api_resource.name,
    config=types.CreateAgentEngineSandboxConfig(display_name="my_custom_sandbox"),
    spec={
        "code_execution_environment": {
            "code_language": language_config,
            "machine_config": machine_config,
        }
    },
)

# Update our sandbox resource name to use this new sandbox
sandbox_resource_name = sandbox_operation.response.name

print("✅ Custom sandbox created successfully!")
print(f"   Language: {language_config}")
print(f"   Machine: {machine_config}")
print(f"   Resource name: {sandbox_resource_name}")

# Initialize Gemini model
model = GenerativeModel("gemini-2.5-flash")

# Ask Gemini to generate code for a calculation
prompt = """
Write Python code to calculate the mean and standard deviation of these numbers:
[23, 45, 67, 89, 12, 34, 56]

Return only the Python code, no explanations.
"""

response = model.generate_content(prompt)
generated_code = response.text.replace("```python", "").replace("```", "").strip()

print("🤖 Gemini generated code:")
print(generated_code)

# Execute the generated code in Agent Engine Sandbox using new pattern
exec_response = client.agent_engines.sandboxes.execute_code(
    name=sandbox_resource_name,
    input_data={"code": generated_code},
)

print("✅ Code executed successfully!")
print("\n📤 Results:")

# Use the new file handling pattern
for output in exec_response.outputs:
    if output.mime_type == "application/json" and output.metadata is None:
        result = json.loads(output.data.decode("utf-8"))
        if result.get("msg_out"):
            print(result.get("msg_out"))
        if result.get("msg_err"):
            print(f"❌ Error: {result.get('msg_err')}")

# Define the code execution as a function for Gemini
def execute_python_code(code: str) -> str:
    """Execute Python code in a secure sandbox.

    Args:
        code: Python code to execute
    Returns:
        The output from code execution
    """
    # Extract code block if wrapped in markdown
    code_match = re.search(r"```python\n(.*?)\n```", code, re.DOTALL)
    if code_match:
        code_to_execute = code_match.group(1)
    else:
        code_to_execute = code

    # Execute in sandbox
    response = client.agent_engines.sandboxes.execute_code(
        name=sandbox_resource_name, input_data={"code": code_to_execute}
    )

    # Parse response using new pattern
    for output in response.outputs:
        if output.mime_type == "application/json" and output.metadata is None:
            result = json.loads(output.data.decode("utf-8"))

            if result.get("msg_err"):
                return f"Error: {result.get('msg_err')}"
            return result.get("msg_out", "Code executed successfully")

    return "Code executed (no output)"

# Create a tool from the function
code_tool = Tool(
    function_declarations=[FunctionDeclaration.from_func(execute_python_code)]
)

# Send a request that will trigger tool use
response = model.generate_content(
    contents=[
        Content(
            role="user",
            parts=[
                Part.from_text(
                    "Calculate the factorial of 10 and check if it's divisible by 100"
                ),
            ],
        )
    ],
    generation_config=GenerationConfig(temperature=0),
    tools=[code_tool],
)

response

# Process the function call from Gemini
function_response_parts = []

for function_call in response.candidates[0].function_calls:
    print(f"🔧 Gemini wants to call: {function_call.name}")
    print(f"📝 Generated code:\n{function_call.args['code']}\n")

    # Execute the code in Agent Engine Sandbox using new pattern
    exec_response = client.agent_engines.sandboxes.execute_code(
        name=sandbox_resource_name, input_data={"code": function_call.args["code"]}
    )

    # Parse response with new pattern
    execution_output = ""
    for output in exec_response.outputs:
        if output.mime_type == "application/json" and output.metadata is None:
            result = json.loads(output.data.decode("utf-8"))

            if result.get("msg_err"):
                execution_output = f"Error: {result.get('msg_err')}"
            else:
                execution_output = result.get("msg_out", "Code executed successfully")

    print(f"✅ Execution result:\n{execution_output}\n")

    # Prepare the function response for Gemini
    function_response_parts.append(
        Part.from_function_response(
            name=function_call.name, response={"result": execution_output}
        )
    )

# Send the function response back to the model
function_response_content = Content(role="function", parts=function_response_parts)

# Get the final response from Gemini
final_response = model.generate_content(
    [
        Content(
            role="user",
            parts=[
                Part.from_text(
                    "Calculate the factorial of 10 and check if it's divisible by 100"
                )
            ],
        ),
        response.candidates[0].content,  # Original function call
        function_response_content,  # Function execution results
    ],
    tools=[code_tool],
)

print("🤖 Gemini's final response:")
print(final_response.text)

# Initialize Claude on Vertex
claude = AnthropicVertex(
    project_id=PROJECT_ID,
    region="us-east5",  # Claude availability varies by region
)

# Ask Claude to generate code
message = claude.messages.create(
    model="claude-sonnet-4@20250514",
    max_tokens=1000,
    messages=[
        {
            "role": "user",
            "content": """Generate Python code to:
        1. Create a list of the first 10 prime numbers
        2. Calculate their sum and average
        3. Create a simple bar chart showing each prime number

        Use matplotlib for the chart. Save the chart as 'primes_chart.png'.
        Return only the Python code.""",
        }
    ],
)

# Extract code from Claude's response
claude_response = message.content[0].text

# Extract code block if wrapped in markdown
code_match = re.search(r"```python\n(.*?)\n```", claude_response, re.DOTALL)
if code_match:
    code_to_execute = code_match.group(1)
else:
    code_to_execute = claude_response

print("🤖 Claude generated code:")
print(code_to_execute)

# Execute Claude's generated code in the Agent Engine Sandbox
exec_response = client.agent_engines.sandboxes.execute_code(
    name=sandbox_resource_name, input_data={"code": code_to_execute}
)

print("✅ Code executed successfully!\n")

# Parse response with new pattern
for output in exec_response.outputs:
    if output.mime_type == "application/json" and output.metadata is None:
        result = json.loads(output.data.decode("utf-8"))

        if result.get("msg_out"):
            print("📤 Execution Output:")
            print(result.get("msg_out"))

        if result.get("msg_err"):
            print("❌ Errors:")
            print(result.get("msg_err"))

# Process all outputs from the sandbox
for output in exec_response.outputs:
    # Handle JSON output (stdout/stderr)
    if output.mime_type == "application/json" and output.metadata is None:
        result = json.loads(output.data.decode("utf-8"))

        if result.get("msg_out"):
            print("📤 Output:")
            print(result.get("msg_out"))

    # Handle generated files (like the chart)
    elif output.metadata and output.metadata.attributes:
        # Extract file name from metadata
        file_name = output.metadata.attributes.get("file_name")
        if isinstance(file_name, bytes):
            file_name = file_name.decode("utf-8")

        print(f"\n📊 Generated file: {file_name}")
        print(f"   MIME type: {output.mime_type}")
        print(f"   Size: {len(output.data)} bytes")

        # If it's an image file, display it
        if file_name.endswith((".png", ".jpg", ".jpeg")):
            print("   Displaying chart...")

            # Decode and display the image
            from io import BytesIO

            import matplotlib.pyplot as plt

            img = plt.imread(BytesIO(output.data))
            fig, ax = plt.subplots(figsize=(8, 6))
            ax.imshow(img)
            ax.axis("off")
            plt.title(f"Generated by Claude: {file_name}")
            plt.show()

            # Optionally save locally
            with open(file_name, "wb") as f:
                f.write(output.data)
            print(f"   ✅ Saved locally as: {file_name}")
        else:
            # For text files, display content
            print(f"   Content: {output.data.decode('utf-8', errors='ignore')}")

# Define the tool schema for Claude
code_execution_tool = {
    "name": "execute_python",
    "description": "Execute Python code in a secure Agent Engine Sandbox",
    "input_schema": {
        "type": "object",
        "properties": {
            "code": {"type": "string", "description": "Python code to execute"}
        },
        "required": ["code"],
    },
}

# Function to handle tool execution with new file handling pattern
def execute_code_tool(code: str) -> str:
    """Execute code when Claude calls the tool."""
    try:
        response = client.agent_engines.sandboxes.execute_code(
            name=sandbox_resource_name, input_data={"code": code}
        )

        # Parse response with new pattern
        for output in response.outputs:
            if output.mime_type == "application/json" and output.metadata is None:
                result = json.loads(output.data.decode("utf-8"))

                if result.get("msg_err"):
                    return f"Error: {result.get('msg_err')}"
                return result.get("msg_out", "Code executed successfully")

        return "Code executed (no output)"
    except Exception as e:
        return f"Execution failed: {e!s}"

# Orchestrate the multi-turn conversation with Claude
message = claude.messages.create(
    model="claude-sonnet-4@20250514",
    max_tokens=1000,
    messages=[
        {"role": "user", "content": "Calculate the 10th Fibonacci number for me."}
    ],
    tools=[code_execution_tool],
)

# Handle tool use in the response
if message.stop_reason == "tool_use":
    tool_results = []

    for content in message.content:
        if content.type == "tool_use":
            print(f"🔧 Claude wants to use tool: {content.name}")
            print(f"📝 With parameters: {content.input}\n")

            # Execute the tool
            if content.name == "execute_python":
                result = execute_code_tool(content.input["code"])
                print(f"✅ Execution result: {result}\n")

                tool_results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": content.id,
                        "content": result,
                    }
                )

    # Send tool results back to Claude for final response
    final_response = claude.messages.create(
        model="claude-sonnet-4@20250514",
        max_tokens=1000,
        messages=[
            {"role": "user", "content": "Calculate the 10th Fibonacci number for me."},
            {"role": "assistant", "content": message.content},
            {"role": "user", "content": tool_results},
        ],
        tools=[code_execution_tool],
    )

    print(f"🤖 Claude's final answer:\n{final_response.content[0].text}")
else:
    # Claude responded without using tools
    print(f"🤖 Claude's response:\n{message.content[0].text}")

# Define the agent with Vertex AI sandbox executor
vertex_agent = LlmAgent(
    model="gemini-2.5-flash",
    name="vertex_code_executor_agent",
    description="An agent that uses Vertex AI Agent Engine Sandbox for code execution",
    instruction="""You are a helpful coding assistant. When asked to perform calculations or data processing:

    1. Write clear, well-commented Python code
    2. Include print statements to show intermediate steps
    3. Use the code executor to run your code
    4. Explain the results in a user-friendly way

    Always ensure your code is complete and executable.
    """,
    code_executor=AgentEngineSandboxCodeExecutor(
        # Use the sandbox we created earlier
        sandbox_resource_name=sandbox_resource_name
    ),
)

print("✅ Agent with AgentEngineSandboxCodeExecutor created!")

# Set up session management
vertex_session_service = InMemorySessionService()
await vertex_session_service.create_session(
    app_name="vertex_code_app", user_id="user1", session_id="session1"
)

artifact_session_service = InMemoryArtifactService()

# Create the runner
vertex_runner = Runner(
    agent=vertex_agent,
    app_name="vertex_code_app",
    session_service=vertex_session_service,
    artifact_service=artifact_session_service,
)

print("✅ Session and runner configured!")

# Run the agent with a computational task
query = "Calculate compound interest for $1000 at 5% annual rate for 10 years"
message = genai_types.Content(role="user", parts=[genai_types.Part(text=query)])

print(f"🙋 User query: {query}\n")
print("=" * 60)

async for event in vertex_runner.run_async(
    user_id="user1", session_id="session1", new_message=message
):
    # Use our helper to parse events
    parse_event(event=event)

builtin_agent = LlmAgent(
    model="gemini-2.5-flash",  # Must use Gemini for BuiltInCodeExecutor
    name="builtin_code_executor_agent",
    description="An agent that uses ADK's built-in code execution",
    instruction="""You are a helpful coding assistant. When asked to perform calculations or data processing:

    1. Write clear, well-commented Python code
    2. Include print statements to show intermediate steps
    3. Use the code execution tool to run your code
    4. Explain the results in a user-friendly way

    Always ensure your code is complete and executable.
    """,
    code_executor=BuiltInCodeExecutor(),  # Much simpler setup!
)

print("✅ Agent with BuiltInCodeExecutor created!")

builtin_session_service = InMemorySessionService()
await builtin_session_service.create_session(
    app_name="builtin_code_app", user_id="user1", session_id="session1"
)

builtin_runner = Runner(
    agent=builtin_agent,
    app_name="builtin_code_app",
    session_service=builtin_session_service,
)

print("✅ Session and runner configured!")

# Run the same query
print(f"🙋 User query: {query}\n")
print("=" * 60)

async for event in builtin_runner.run_async(
    user_id="user1", session_id="session1", new_message=message
):
    # Our helper handles both event types
    parse_event(event)

class DataAnalysisResult(BaseModel):
    """Structured output for data analysis results."""

    total_sales: float = Field(description="Total sales amount")
    average_sales: float = Field(description="Average sales per product")
    top_product: str = Field(description="Product with highest sales")
    insights: str = Field(description="Key insights from the analysis")

data_analyst = LlmAgent(
    model="gemini-2.5-flash",
    name="data_analyst",
    description="Expert data analyst for sales and business metrics",
    instruction="""You are an expert data analyst. When given data:

    1. First, load and explore the data structure
    2. Calculate key metrics (totals, averages, trends) using code executor
    3. Identify top performers and outliers
    4. Generate actionable insights

    Always use pandas for data manipulation and include clear print statements.
    Format numbers nicely (e.g., currency with commas).
    """,
    code_executor=AgentEngineSandboxCodeExecutor(
        sandbox_resource_name=sandbox_resource_name
    ),
    output_key="analysis_result",  # Store result in session state
)

print("✅ Data Analyst agent created!")

# Initialize session with all services
analyst_session_service = InMemorySessionService()
analyst_memory_service = InMemoryMemoryService()
analyst_artifact_service = InMemoryArtifactService()

analyst_session = await analyst_session_service.create_session(
    app_name="sales_analysis",
    user_id="analyst_001",
    session_id="analysis_123",
    state={},  # Empty initial state
)

# Create runner
analyst_runner = Runner(
    agent=data_analyst,
    app_name="sales_analysis",
    session_service=analyst_session_service,
    memory_service=analyst_memory_service,
    artifact_service=analyst_artifact_service,
)

print("✅ Data Analyst runner configured!")

# Prepare the analysis request with CSV data
analysis_request = genai_types.Content(
    role="user",
    parts=[
        genai_types.Part(
            text="""
Analyze this sales data and provide insights:

Product,Sales,Units,Region
Widget A,5000,100,North
Widget B,7500,150,South
Widget C,3000,60,East
Widget D,9000,180,West
Widget E,6500,130,North

Calculate:
1. Total and average sales
2. Best performing product
3. Sales per unit for each product
4. Regional performance summary
"""
        )
    ],
)

print("🙋 User query: Sales data analysis\n")
print("=" * 60)

# Run the analysis
async for event in analyst_runner.run_async(
    user_id="analyst_001", session_id="analysis_123", new_message=analysis_request
):
    parse_event(event)

# List all sandboxes in the Agent Engine
sandboxes = client.agent_engines.sandboxes.list(name=agent_engine.api_resource.name)

print(f"✅ Found {len(sandboxes)} sandbox(es)\n")
print("=" * 60)

for i, sandbox in enumerate(sandboxes, 1):
    print(f"\n📦 Sandbox {i}:")
    print(f"   Display name: {sandbox.display_name}")
    print(f"   Resource name: {sandbox.name}")
    print(f"   State: {sandbox.state}")
    print(f"   Created: {sandbox.create_time}")
    if hasattr(sandbox, "expire_time") and sandbox.expire_time:
        print(f"   Expires: {sandbox.expire_time}")

# Get detailed information about a specific sandbox
if sandboxes:
    sandbox_name = sandboxes[0].name
    sandbox = client.agent_engines.sandboxes.get(name=sandbox_name)

    print("✅ Sandbox details retrieved!\n")
    print(f"📦 Sandbox: {sandbox.display_name}")
    print(f"   State: {sandbox.state}")
    print(f"   Created: {sandbox.create_time}")
    print(f"   Spec: {sandbox.spec}")
else:
    print("⚠️ No sandboxes found to inspect")

# Delete a specific sandbox (with error handling)
if sandboxes and len(sandboxes) > 1:  # Only if we have more than one
    try:
        # Delete the first sandbox (not the one we're actively using)
        delete_operation = client.agent_engines.sandboxes.delete(name=sandboxes[0].name)

        if delete_operation.done:
            print("✅ Sandbox deleted successfully!")
            print(f"   Deleted: {sandboxes[0].display_name}")
        else:
            print("⏳ Deletion in progress...")
            print(f"   Resource: {sandboxes[0].display_name}")
    except Exception as e:
        print(f"⚠️ Error during deletion: {e!s}")
else:
    print("ℹ️ Skipping deletion (only one sandbox or none available)")

# Define the code to read and write text files
my_code = """
with open("input.txt", "r") as input_file:
    with open("output.txt", "w") as output_file:
        for line in input_file:
            # Echo each line to stdout for visibility
            print(f"Processing: {line.strip()}")
            # Write to output file
            output_file.write(line)
"""

# Prepare input data with a file
input_data = {
    "code": my_code,
    "files": [
        {
            "name": "input.txt",
            "content": b"Hello, Agent Engine Sandbox!\nThis is a test file.\nFile I/O is working!",
        }
    ],
}

# Execute the code
response = client.agent_engines.sandboxes.execute_code(
    name=sandbox_resource_name, input_data=input_data
)

print("✅ Code executed with file I/O!")

# Process the response outputs
for output in response.outputs:
    # Check if this is JSON output (stdout/stderr)
    if output.mime_type == "application/json" and output.metadata is None:
        # Decode the JSON response
        result = json.loads(output.data.decode("utf-8"))

        # Display stdout
        if result.get("msg_out"):
            print("📤 Standard Output:")
            print(result.get("msg_out"))

        # Display stderr if any
        if result.get("msg_err"):
            print("❌ Errors:")
            print(result.get("msg_err"))

    # Check if this is a generated file
    elif output.metadata and output.metadata.attributes:
        # Extract the file name from metadata
        file_name = output.metadata.attributes.get("file_name")
        if isinstance(file_name, bytes):
            file_name = file_name.decode("utf-8")

        print(f"\n📁 Generated File: {file_name}")
        print(f"   MIME Type: {output.mime_type}")
        print(f"   Size: {len(output.data)} bytes")
        print(
            f"   Content preview: {output.data[:100].decode('utf-8', errors='ignore')}..."
        )

# Code to generate a matplotlib chart
chart_code = """
import matplotlib.pyplot as plt

# Create data
x = [1, 2, 3, 4, 5]
y = [2, 4, 6, 8, 10]

# Create the plot
plt.figure(figsize=(8, 6))
plt.plot(x, y, marker='o', linewidth=2, markersize=8)
plt.xlabel('X-axis')
plt.ylabel('Y-axis')
plt.title('Simple Line Plot')
plt.grid(True, alpha=0.3)

# Save the chart
plt.savefig('chart_out.png', dpi=150, bbox_inches='tight')
print("Chart saved to 'chart_out.png'")
"""

# Execute the code (no input files needed this time)
response = client.agent_engines.sandboxes.execute_code(
    name=sandbox_resource_name, input_data={"code": chart_code}
)

print("✅ Chart generation code executed!")

# Process outputs to find and display the image
for output in response.outputs:
    # Handle JSON output (stdout/stderr)
    if output.mime_type == "application/json" and output.metadata is None:
        result = json.loads(output.data.decode("utf-8"))
        if result.get("msg_out"):
            print("📤 Output:")
            print(result.get("msg_out"))

    # Handle image files
    elif output.metadata and output.metadata.attributes:
        file_name = output.metadata.attributes.get("file_name")
        if isinstance(file_name, bytes):
            file_name = file_name.decode("utf-8")

        print(f"\n📊 Generated Image: {file_name}")
        print(f"   MIME Type: {output.mime_type}")
        print(f"   Size: {len(output.data)} bytes")

        # Display the image if it's a PNG/JPG
        if file_name.endswith((".png", ".jpg", ".jpeg")):
            # Decode the binary data and display
            img = plt.imread(BytesIO(output.data))
            fig, ax = plt.subplots(figsize=(8, 6))
            ax.imshow(img)
            ax.axis("off")
            plt.title(f"Retrieved: {file_name}")
            plt.show()

            # Optionally save locally
            with open(file_name, "wb") as f:
                f.write(output.data)
            print(f"   ✅ Saved locally as: {file_name}")

# Clean up the Agent Engine and all child resources
delete_agent_engine = True  # Set to False to keep resources

if delete_agent_engine:
    try:
        # Using force=True will delete all child sandboxes automatically
        agent_engine.delete(force=True)
        print("✅ Agent Engine and all sandboxes deleted successfully!")
        print("   All resources have been cleaned up.")
    except Exception as e:
        print(f"⚠️ Error during cleanup: {e!s}")
else:
    print("ℹ️ Keeping Agent Engine resources (delete_agent_engine = False)")