# Copyright 2026 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.

PROJECT_ID = "YOUR_PROJECT_ID"  # @param {type: "string"}

%pip install --upgrade google-cloud-aiplatform[agent_engines,adk]>=1.144

from google.colab import auth

auth.authenticate_user()

import vertexai

LOCATION = "us-central1"  # @param {type: "string"}
ENDPOINT = f"https://{LOCATION}-aiplatform.googleapis.com"

client = vertexai.Client(
    project=PROJECT_ID,
    location=LOCATION,
)

!mkdir -p echo_agent

%%writefile requirements.txt
google-cloud-aiplatform[agent_engines,adk]
fastapi
uvicorn
pydantic

%%writefile Dockerfile
FROM python:3.13-alpine
WORKDIR /app

RUN adduser --disabled-password --gecos "" myuser

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt


COPY --chown=myuser:myuser . .

ENV PATH="/home/myuser/.local/bin:$PATH"

USER myuser

CMD ["sh", "-c", "uvicorn main:app --host 0.0.0.0 --port $PORT"]

%%writefile echo_agent/__init__.py
from . import agent

%%writefile echo_agent/agent.py
import asyncio
import logging

from typing import Any, Dict, AsyncIterator

class EchoAgent:
    """
    A simple echo agent demonstrating bidirectional WebSocket capabilities.
    """
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.logger.info("Echo Agent ready!")

    async def bidi_stream_query(
        self,
        queue: asyncio.Queue
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        Bidirectional streaming for continuous conversation.
        """
        self.logger.info("Bidi session started")

        while True:
            # Wait for message from the queue
            message = await queue.get()
            user_input = message.get("message", "")

            # Check for exit command
            if user_input.lower() in ("exit", "quit"):
                yield {"output": "Goodbye!"}
                break

            # Echo back the input
            # In a real agent, this is where you'd process the input
            yield {"output": f"Echo: {user_input}"}

# Create an instance of our echo agent
root_agent = EchoAgent()

%%writefile main.py
import os
import asyncio
import logging
import uvicorn
from fastapi import FastAPI, WebSocket

from echo_agent.agent import root_agent

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("bidi_server")
app = FastAPI()

@app.websocket("/bidi_echo")
async def bidi_handler(websocket: WebSocket):
    await websocket.accept()

    initial_data = await websocket.receive_json(mode="binary")
    logger.info(f"Inbound initial handshake message: {initial_data}")

    queue = asyncio.Queue()
    # Seed the queue if the initial frame contains a starting payload
    if "input" in initial_data:
        queue.put_nowait(initial_data["input"])

    async def receive_messages():
        while True:
            try:
                data = await websocket.receive_json(mode="binary")
                logger.info(f"Inbound bidi message: {data}")

                queue.put_nowait(data.get("input", {}))
            except Exception as e:
                logger.info(f"WebSocket connection closed or receive error: {e}")
                break

    async def send_messages():
        try:
            async for response in root_agent.bidi_stream_query(queue):
                await websocket.send_json({"output": response}, mode="binary")
        except Exception as e:
            logger.error(f"Error sending bidi response: {e}")
    # Run WebSocket communication loops
    await asyncio.gather(receive_messages(), send_messages())

# You can add more FastAPI routes or configurations below if needed
# Example:
# @app.get("/hello")
# async def read_root():
#     return {"Hello": "World"}

if __name__ == "__main__":
    # Use the PORT environment variable provided by Cloud Run, defaulting to 8080
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))

remote_echo_agent = client.agent_engines.create(
    config={
        "display_name": "Bidi Echo test agent",
        "description": "bidi stream testing with echo agent",
        "source_packages": [
            # The files in the current directory to upload. You can also use "."
            "echo_agent",
            "Dockerfile",
            "main.py",
            "requirements.txt",
        ],
        "image_spec": {},  # tells AgentRuntime to use the Dockerfile
        "agent_framework": "google-adk",  # For usage through the console / UI
        "env_vars": {"GOOGLE_GENAI_USE_VERTEXAI": "1"},
        "resource_limits": {"cpu": "2", "memory": "8Gi"},
        # Explicitly set max_instances to stay within quota
        "max_instances": 5,
        # You can also adjust min_instances if needed, default is 1
        # "min_instances": 1,
        # # Cloud Storage bucket for staging
        # "staging_bucket": BUCKET_URI,
    },
)
remote_echo_agent

# Put the echo agent's reasoning engine id from the last step here
ECHO_ENGINE_ID = "1266316887558455296"  # @param {type: "string"}

import asyncio
import json

import google.auth
import websockets
from google.auth.transport.requests import Request

creds, _ = google.auth.default()
if not creds.valid:
    creds.refresh(Request())

AGENT_PATH = "bidi_echo"  # This should match the API server's handler path
url = f"wss://{ENDPOINT.removeprefix('https://')}/reasoningEngines/ws/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{ECHO_ENGINE_ID}/api/{AGENT_PATH}"
print(url)


async def test_websocket():
    headers = {
        "Authorization": f"Bearer {creds.token}",
        "Content-Type": "application/octet-stream",
    }

    messages = ["first request", "second request", "almost done!", "exit"]

    try:
        async with websockets.connect(url, additional_headers=headers) as websocket:
            print(f"Connected to {url}\n")

            for msg in messages:
                request_dict = {"input": {"message": msg}}
                print(f"Sending: {request_dict}")
                await websocket.send(json.dumps(request_dict))

                # Listen for response
                try:
                    response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    print(f"Received: {response}\n")
                except asyncio.TimeoutError:
                    print("No message received within the timeout period.\n")

    except Exception as e:
        print(f"An error occurred: {e}")
        raise e


await test_websocket()

!mkdir -p capital_agent

%%writefile requirements.txt
google-cloud-aiplatform[agent_engines,adk]
fastapi
uvicorn
pydantic

%%writefile Dockerfile
FROM python:3.13-alpine
WORKDIR /app

RUN adduser --disabled-password --gecos "" myuser

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt


COPY --chown=myuser:myuser . .

ENV PATH="/home/myuser/.local/bin:$PATH"

USER myuser

CMD ["sh", "-c", "uvicorn main:app --host 0.0.0.0 --port $PORT"]

%%writefile capital_agent/__init__.py
from . import agent

%%writefile capital_agent/agent.py
from google.adk.agents import Agent

# Define a tool function
def get_capital_city(country: str) -> str:
  """Retrieves the capital city for a given country."""
  # Replace with actual logic (e.g., API call, database lookup)
  capitals = {"france": "Paris", "japan": "Tokyo", "canada": "Ottawa"}
  return capitals.get(country.lower(), f"Sorry, I don't know the capital of {country}.")

root_agent = Agent(
    model="gemini-live-2.5-flash-native-audio",  # Specify the Live API model
    name="capital_agent",
    description="Answers user questions about the capital city of a given country using Live API.",
    instruction="""You are an agent that provides the capital city of a country.
             When a user asks for the capital of a country:
             1. Identify the country name from the user's query.
             2. Use the `get_capital_city` tool to find the capital.
             3. Respond clearly to the user, stating the capital city.
             Example Query: "What's the capital of France?"
             Example Response: "The capital of France is Paris."
        """,
    tools=[get_capital_city]
)

%%writefile main.py
import os

import uvicorn
from fastapi import FastAPI
from google.adk.cli.fast_api import get_fast_api_app

# Get the directory where main.py is located
AGENT_DIR = os.path.dirname(os.path.abspath(__file__))
# Example session service URI (e.g., SQLite)
SESSION_SERVICE_URI = ""
# Example allowed origins for CORS
ALLOWED_ORIGINS = ["http://localhost", "http://localhost:8080", "*"]
# Set web=True if you intend to serve a web interface, False otherwise
SERVE_WEB_INTERFACE = True

# Call the function to get the FastAPI app instance
# Ensure the agent directory name ('capital_agent') matches your agent folder
app: FastAPI = get_fast_api_app(
    agents_dir=AGENT_DIR,
    session_service_uri=SESSION_SERVICE_URI,
    allow_origins=ALLOWED_ORIGINS,
    web=SERVE_WEB_INTERFACE,
)

# You can add more FastAPI routes or configurations below if needed
# Example:
# @app.get("/hello")
# async def read_root():
#     return {"Hello": "World"}

if __name__ == "__main__":
    # Use the PORT environment variable provided by Cloud Run, defaulting to 8080
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))

remote_live_agent = client.agent_engines.create(
    config={
        "display_name": "Bidi Live API test agent",
        "description": "bidi stream testing with live agent",
        "source_packages": [
            # The files in the current directory to upload.
            "capital_agent",
            "Dockerfile",
            "main.py",
            "requirements.txt",
        ],
        "image_spec": {},  # tells AgentRuntime to use the Dockerfile
        "agent_framework": "google-adk",  # For usage through the console / UI
        "env_vars": {"GOOGLE_GENAI_USE_VERTEXAI": "1"},
        "resource_limits": {"cpu": "2", "memory": "8Gi"},
        # Explicitly set max_instances to stay within quota
        "max_instances": 5,
        # You can also adjust min_instances if needed, default is 1
        # "min_instances": 1,
    },
)
remote_live_agent

# Put the live agent's reasoning engine id from the last step here
LIVE_ENGINE_ID = "7678296516062085120"  # @param {type: "string"}

APP_NAME = "capital_agent"
USER_ID = "demo_user"

# Put the session id you want to create
SESSION_ID = "demo_session"  # @param {type: "string"}

import google.auth
import requests
from google.auth.transport.requests import Request

creds, _ = google.auth.default()
if not creds.valid:
    creds.refresh(Request())

base_ingress_url = f"{ENDPOINT}/reasoningEngines/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api"
create_session_url = (
    f"{base_ingress_url}/apps/{APP_NAME}/users/{USER_ID}/sessions/{SESSION_ID}"
)

headers = {"Authorization": f"Bearer {creds.token}", "Content-Type": "application/json"}

response = requests.post(create_session_url, headers=headers)
response.json()

import google.auth
import requests
from google.auth.transport.requests import Request

creds, _ = google.auth.default()
if not creds.valid:
    creds.refresh(Request())

base_ingress_url = f"{ENDPOINT}/reasoningEngines/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api"
list_sessions_url = f"{base_ingress_url}/apps/{APP_NAME}/users/{USER_ID}/sessions"

headers = {"Authorization": f"Bearer {creds.token}", "Content-Type": "application/json"}

response = requests.get(list_sessions_url, headers=headers)
response.json()

LIVE_QUERY_TEXT = "What is the capital of France?"  # @param {type: "string"}

import base64
import time
import wave

import google.auth
from IPython.display import Audio, display
from google.auth.transport.requests import Request

creds, _ = google.auth.default()
if not creds.valid:
    creds.refresh(Request())

# Parameters for the audio
sample_rate = 44100  # samples per second

url = f"wss://{ENDPOINT.removeprefix('https://')}/reasoningEngines/ws/internal/projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{LIVE_ENGINE_ID}/api/run_live?app_name={APP_NAME}&user_id={USER_ID}&session_id={SESSION_ID}"


def save_wav_file(filename, pcm_data, channels=1, sampwidth=2, framerate=24000):
    """Saves raw PCM data to a WAV file."""
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(sampwidth)  # 2 bytes for 16-bit audio
        wf.setframerate(framerate)
        wf.writeframes(pcm_data)
    print(f"Audio saved to {filename}")


async def test_websocket():
    headers = {
        "Authorization": f"Bearer {creds.token}",
    }

    try:
        async with websockets.connect(url, additional_headers=headers) as websocket:
            print(f"Connected to {url}")
            test_message = {
                "content": {"role": "user", "parts": [{"text": LIVE_QUERY_TEXT}]}
            }
            print(f"Sending: {test_message}")
            await websocket.send(json.dumps(test_message))
            print(f"Sent: {test_message}")

            # Listen for responses
            print("--- Receiving Responses ---")
            audio_buffer = b""
            text_response = ""
            try:
                while True:
                    response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    print(f"Received: {response}")
                    try:
                        data = json.loads(response)

                        if "content" in data and data["content"]["role"] == "model":
                            for part in data["content"].get("parts", []):
                                if "text" in part:
                                    text_response += part["text"]
                                    print(
                                        f"====== MODEL TEXT PART: {part['text']} ======"
                                    )
                                if "inlineData" in part:
                                    mime_type = part["inlineData"].get("mimeType")
                                    if mime_type == "audio/pcm":
                                        b64_data = part["inlineData"].get("data")
                                        if b64_data:
                                            try:
                                                print(type(b64_data))
                                                decoded_chunk = (
                                                    base64.urlsafe_b64decode(b64_data)
                                                )
                                                audio_buffer += decoded_chunk
                                                print(
                                                    f"Decoded and appended {len(decoded_chunk)} audio bytes."
                                                )
                                            except base64.binascii.Error as e:
                                                print(b64_data)
                                                print(
                                                    f"Base64 Decode Error: {e} - on chunk length {len(b64_data)}"
                                                )

                        if data.get("turnComplete"):
                            print("Turn complete received.")
                            if text_response:
                                print(
                                    f"====== FINAL TEXT RESPONSE: {text_response} ======"
                                )
                            if audio_buffer:
                                print("Saving audio buffer")
                                filename = f"response_{int(time.time())}.wav"
                                save_wav_file(filename, audio_buffer)
                                display(Audio(filename=filename, rate=sample_rate))
                                audio_buffer = b""
                                break
                    except json.JSONDecodeError:
                        print(f"Received non-JSON message: {response}")
            except asyncio.TimeoutError:
                print("No message received within the timeout period.")
            except websockets.exceptions.ConnectionClosedOK:
                print("Connection closed normally.")
            except websockets.exceptions.ConnectionClosedError as e:
                print(f"Connection closed with error: {e}")

    except Exception as e:
        print(f"An error occurred: {e}")
        raise e


await test_websocket()

remote_echo_agent.delete(force=True)
remote_live_agent.delete(force=True)