# 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
#
#     http://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.

# # --- Setup Instructions ---
# # 1. Install the ADK package:
# !pip install google-adk
# # Make sure to restart kernel if using colab/jupyter notebooks

# # 2. Set up your Gemini API Key:
# #    - Get a key from Google AI Studio: https://aistudio.google.com/app/apikey
# #    - Set it as an environment variable:
# import os
# os.environ["GOOGLE_API_KEY"] = "YOUR_API_KEY_HERE" # <--- REPLACE with your actual key
# # Or learn about other authentication methods (like Agent Platform):
# # https://adk.dev/agents/models/

# ADK Imports
from google.adk.agents import LlmAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.runners import InMemoryRunner  # Use InMemoryRunner
from google.genai import types  # For types.Content
from typing import Optional

# Define the model - Use the specific model name requested
GEMINI_2_FLASH = "gemini-2.0-flash"


# --- 1. Define the Callback Function ---
def check_if_agent_should_run(
    callback_context: CallbackContext,
) -> Optional[types.Content]:
    """
    Logs entry and checks 'skip_llm_agent' in session state.
    If True, returns Content to skip the agent's execution.
    If False or not present, returns None to allow execution.
    """
    agent_name = callback_context.agent_name
    invocation_id = callback_context.invocation_id
    current_state = callback_context.state.to_dict()

    print(f"\n[Callback] Entering agent: {agent_name} (Inv: {invocation_id})")
    print(f"[Callback] Current State: {current_state}")

    # Check the condition in session state dictionary
    if current_state.get("skip_llm_agent", False):
        print(
            f"[Callback] State condition 'skip_llm_agent=True' met: Skipping agent {agent_name}."
        )
        # Return Content to skip the agent's run
        return types.Content(
            parts=[
                types.Part(
                    text=f"Agent {agent_name} skipped by before_agent_callback due to state."
                )
            ],
            role="model",  # Assign model role to the overriding response
        )
    else:
        print(
            f"[Callback] State condition not met: Proceeding with agent {agent_name}."
        )
        # Return None to allow the LlmAgent's normal execution
        return None


# --- 2. Setup Agent with Callback ---
llm_agent_with_before_cb = LlmAgent(
    name="MyControlledAgent",
    model=GEMINI_2_FLASH,
    instruction="You are a concise assistant.",
    description="An LLM agent demonstrating stateful before_agent_callback",
    before_agent_callback=check_if_agent_should_run,  # Assign the callback
)


# --- 3. Setup Runner and Sessions using InMemoryRunner ---
async def main():
    app_name = "before_agent_demo"
    user_id = "test_user"
    session_id_run = "session_will_run"
    session_id_skip = "session_will_skip"

    # Use InMemoryRunner - it includes InMemorySessionService
    runner = InMemoryRunner(agent=llm_agent_with_before_cb, app_name=app_name)
    # Get the bundled session service to create sessions
    session_service = runner.session_service

    # Create session 1: Agent will run (default empty state)
    await session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id_run,
        # No initial state means 'skip_llm_agent' will be False in the callback check
    )

    # Create session 2: Agent will be skipped (state has skip_llm_agent=True)
    await session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id_skip,
        state={"skip_llm_agent": True},  # Set the state flag here
    )

    # --- Scenario 1: Run where callback allows agent execution ---
    print(
        "\n"
        + "=" * 20
        + f" SCENARIO 1: Running Agent on Session '{session_id_run}' (Should Proceed) "
        + "=" * 20
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id_run,
        new_message=types.Content(
            role="user", parts=[types.Part(text="Hello, please respond.")]
        ),
    ):
        # Print final output (either from LLM or callback override)
        if event.is_final_response() and event.content:
            print(
                f"Final Output: [{event.author}] {event.content.parts[0].text.strip()}"
            )
        elif event.is_error():
            print(f"Error Event: {event.error_details}")

    # --- Scenario 2: Run where callback intercepts and skips agent ---
    print(
        "\n"
        + "=" * 20
        + f" SCENARIO 2: Running Agent on Session '{session_id_skip}' (Should Skip) "
        + "=" * 20
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id_skip,
        new_message=types.Content(
            role="user", parts=[types.Part(text="This message won't reach the LLM.")]
        ),
    ):
        # Print final output (either from LLM or callback override)
        if event.is_final_response() and event.content:
            print(
                f"Final Output: [{event.author}] {event.content.parts[0].text.strip()}"
            )
        elif event.is_error():
            print(f"Error Event: {event.error_details}")


# --- 4. Execute ---
# In a Python script:
# import asyncio
# if __name__ == "__main__":
#     # Make sure GOOGLE_API_KEY environment variable is set if not using Agent Platform auth
#     # Or ensure Application Default Credentials (ADC) are configured for Agent Platform
#     asyncio.run(main())

# In a Jupyter Notebook or similar environment:
await main()