#!/usr/bin/env python3
"""
memory_setup.py — Initialize Agent Engine instance with Sessions + Memory Bank.
Project: propane-will-491900-m5 | Region: us-central1
SDK: google-cloud-aiplatform >= 1.111.0

COST NOTE: Sessions + Memory Bank are metered since Jan 28, 2026.
Trim session histories. Only persist high-value facts to long-term memory.
"""

import vertexai
from vertexai import Client
from google.adk.memory import VertexAiMemoryBankService
from google.adk.sessions import VertexAiSessionService

PROJECT_ID = "propane-will-491900-m5"
LOCATION = "us-central1"


def create_agent_engine() -> str:
    """Create an Agent Engine instance (backing store for Sessions + Memory Bank)."""
    client = Client(project=PROJECT_ID, location=LOCATION)
    agent_engine = client.agent_engines.create()
    resource_name = agent_engine.api_resource.name
    agent_engine_id = resource_name.split("/")[-1]
    print(f"Agent Engine created: {resource_name}")
    print(f"Agent Engine ID: {agent_engine_id}")
    print("Store this in Secret Manager or env var: AGENT_ENGINE_ID")
    return agent_engine_id


def get_services(agent_engine_id: str):
    """Return configured session and memory services for use with ADK Runner."""
    memory_service = VertexAiMemoryBankService(
        project=PROJECT_ID,
        location=LOCATION,
        agent_engine_id=agent_engine_id,
    )
    session_service = VertexAiSessionService(
        project_id=PROJECT_ID,
        location=LOCATION,
        agent_engine_id=agent_engine_id,
    )
    return session_service, memory_service


if __name__ == "__main__":
    agent_engine_id = create_agent_engine()
    session_svc, memory_svc = get_services(agent_engine_id)
    print("Services ready.")
    print(f"session_service: {session_svc}")
    print(f"memory_service:  {memory_svc}")