# 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 google-cloud-aiplatform

import sys

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

    auth.authenticate_user()

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")

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

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

import datetime
import os
import uuid
import warnings

warnings.filterwarnings("ignore")

# Import class-based types for Memory Bank
from vertexai import types

# Basic configuration types
MemoryBankConfig = types.ReasoningEngineContextSpecMemoryBankConfig
SimilaritySearchConfig = (
    types.ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig
)
GenerationConfig = types.ReasoningEngineContextSpecMemoryBankConfigGenerationConfig

# Advanced configuration types
TtlConfig = types.ReasoningEngineContextSpecMemoryBankConfigTtlConfig
GranularTtlConfig = (
    types.ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig
)
CustomizationConfig = types.MemoryBankCustomizationConfig
MemoryTopic = types.MemoryBankCustomizationConfigMemoryTopic
ManagedMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic
CustomMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic
GenerateMemoriesExample = types.MemoryBankCustomizationConfigGenerateMemoriesExample
ConversationSource = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource
)
ConversationSourceEvent = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent
)
ExampleGeneratedMemory = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory
)
ManagedTopicEnum = types.ManagedTopicEnum

print("✅ Libraries imported successfully!")

def display_generated_memories(operation, client, title="Generated Memories"):
    """Display memories from a generation operation with consistent formatting.

    Args:
        operation: The result from client.agent_engines.memories.generate()
        client: The Vertex AI client instance
        title: Title to display above the memories
    """
    if operation.response and operation.response.generated_memories:
        print(f"\n✅ {title}: {len(operation.response.generated_memories)}\n")

        for i, gen_memory in enumerate(operation.response.generated_memories, 1):
            if gen_memory.action != "DELETED" and gen_memory.memory:
                try:
                    full_memory = client.agent_engines.memories.get(
                        name=gen_memory.memory.name
                    )
                    action_icon = "🆕" if gen_memory.action == "CREATED" else "🔄"
                    print(f"   {action_icon} {i}. {full_memory.fact}")
                except Exception as e:
                    print(f"   ⚠️ Could not retrieve memory: {e}")
    else:
        print(f"\n📭 No {title.lower()} found")


print("✅ Helper function defined!")

print("🧠 Creating Memory Bank configuration for hotel concierge...\n")

basic_memory_config = MemoryBankConfig(
    # Which embedding model to use for similarity search
    similarity_search_config=SimilaritySearchConfig(
        embedding_model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005"
    ),
    # Which LLM to use for extracting memories from conversations
    generation_config=GenerationConfig(
        model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    ),
)

print("✅ Memory Bank configuration created!")
print("   Embedding model: text-embedding-005")
print("   Generation model: gemini-2.5-flash")

print("\n🛠️ Creating Agent Engine with Memory Bank...\n")
print("⏳ This provisions the backend infrastructure for guest memory storage...")

agent_engine = client.agent_engines.create(
    config={"context_spec": {"memory_bank_config": basic_memory_config}}
)

agent_engine_name = agent_engine.api_resource.name

print("\n✅ Agent Engine created successfully!")
print(f"   Resource Name: {agent_engine_name}")

print("💬 Creating a session for guest check-in...\n")

# Generate a unique guest identifier
guest_id = "guest_emma_" + str(uuid.uuid4())[:4]

# Create a session for this guest
session = client.agent_engines.sessions.create(
    name=agent_engine_name,
    user_id=guest_id,
    config={"display_name": f"Check-in conversation for {guest_id}"},
)

session_name = session.response.name

print("✅ Session created successfully!")
print(f"   Session: {session_name}")
print(f"   Guest ID: {guest_id}")

# A realistic hotel check-in conversation
guest_conversation = [
    {
        "role": "user",
        "message": "Hi! I'm Emma Chen. I have a reservation for a deluxe room for three nights.",
    },
    {
        "role": "model",
        "message": "Welcome to our hotel, Ms. Chen! I have your reservation right here. Let me get you checked in.",
    },
    {
        "role": "user",
        "message": "Great! Just so you know, I'm vegetarian and lactose intolerant, so please note that for room service and the breakfast buffet.",
    },
    {
        "role": "model",
        "message": "Absolutely, I've noted your dietary preferences. Our kitchen will be informed, and we have excellent vegetarian options.",
    },
    {
        "role": "user",
        "message": "Thank you! Also, I prefer a high floor room, away from the elevator. I'm a light sleeper, so somewhere quiet would be ideal.",
    },
    {
        "role": "model",
        "message": "Of course! I'll assign you a corner room on the 12th floor. It's one of our quietest locations.",
    },
    {
        "role": "user",
        "message": "Perfect! One more thing - can I get extra pillows and a white noise machine? And I prefer the room temperature around 68°F.",
    },
    {
        "role": "model",
        "message": "Certainly! I'll have housekeeping bring extra pillows and a white noise machine to your room, and I'll set the thermostat to 68°F before your arrival.",
    },
]

print("✅ Guest conversation defined!")
print(f"   {len(guest_conversation)} conversation turns")

print("⬆️ Adding conversation to session...\n")

invocation_id = 0

for turn in guest_conversation:
    client.agent_engines.sessions.events.append(
        name=session_name,
        author=guest_id,  # Required by Sessions
        invocation_id=str(invocation_id),  # Required by Sessions
        timestamp=datetime.datetime.now(
            tz=datetime.timezone.utc
        ),  # Required by Sessions
        config={
            "content": {"role": turn["role"], "parts": [{"text": turn["message"]}]}
        },
    )

    invocation_id += 1
    icon = "👤" if turn["role"] == "user" else "🤖"
    print(f"{icon} {turn['message']}")

print("\n✅ Conversation successfully added to session!")

print("🧠 Analyzing conversation and extracting memories...\n")
print(
    "⏳ The AI is reading the conversation and identifying important facts to remember..."
)

# Generate memories from the session
operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    vertex_session_source={"session": session_name},
    config={"wait_for_completion": True},
)

print("\n✅ Memory extraction complete!")

# Display the generated memories using our helper function
display_generated_memories(operation, client, "Guest Preferences Extracted")

print(f"📚 Retrieving all memories for guest: {guest_id}...\n")

# Simple retrieval - get all memories for this guest
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": guest_id}
)

all_memories = list(results)

print(f"✅ Found {len(all_memories)} memories for this guest:\n")

for i, retrieved_memory in enumerate(all_memories, 1):
    print(f"   {i}. {retrieved_memory.memory.fact}")

print("=" * 80)
print("SCENARIO: Guest Returns for Second Stay")
print("=" * 80)
print("\n📅 Two weeks later, Emma returns to the hotel...\n")

# Simulate a return visit conversation
print("👤 Hi! I'm Emma Chen, checking in again.")
print("🤖 Welcome back, Ms. Chen! Let me pull up your preferences...\n")

# Retrieve Emma's memories
print("🔍 Retrieving guest preferences from previous stay...")
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": guest_id}
)

memories = list(results)

print(f"✅ Found {len(memories)} preferences on file\n")

# Display how we'll use these memories
print("🤖 Perfect! Based on your previous stay, I have:")
print()

for memory in memories:
    fact = memory.memory.fact

    # Provide contextual responses based on memory content
    if "vegetarian" in fact.lower() or "lactose" in fact.lower():
        print(f"   🥗 Dietary: {fact}")
        print("      → I've already notified the kitchen and breakfast staff")
    elif (
        "floor" in fact.lower() or "quiet" in fact.lower() or "elevator" in fact.lower()
    ):
        print(f"   🏨 Room: {fact}")
        print("      → I've assigned you the same corner room on floor 12")
    elif "temperature" in fact.lower() or "68" in fact:
        print(f"   🌡️  Climate: {fact}")
        print("      → Thermostat pre-set before your arrival")
    elif "pillow" in fact.lower() or "white noise" in fact.lower():
        print(f"   🛏️  Amenities: {fact}")
        print("      → Already placed in your room")
    else:
        print(f"   ℹ️  {fact}")

    print()

print("📞 Emma calls the front desk with some questions...\n")

# Additional conversation turns
additional_conversation = [
    {
        "role": "user",
        "message": "Hi, I'd like to book a spa treatment tomorrow. Do you have deep tissue massage? I prefer morning appointments, around 9 AM.",
    },
    {
        "role": "model",
        "message": "Absolutely! I can book you a deep tissue massage at 9 AM tomorrow. I've noted your preference for morning spa appointments.",
    },
    {
        "role": "user",
        "message": "Great! Also, what vegetarian options do you have for breakfast? I need dairy-free choices because of my lactose intolerance.",
    },
    {
        "role": "model",
        "message": "We have several vegetarian and dairy-free options: oat milk smoothie bowls, avocado toast, fresh fruit platters, and coconut yogurt parfaits.",
    },
]

# Add these turns to the same session
invocation_id_counter = len(guest_conversation)

for turn in additional_conversation:
    client.agent_engines.sessions.events.append(
        name=session_name,
        author=guest_id,
        invocation_id=str(invocation_id_counter),
        timestamp=datetime.datetime.now(tz=datetime.timezone.utc),
        config={
            "content": {"role": turn["role"], "parts": [{"text": turn["message"]}]}
        },
    )

    invocation_id_counter += 1
    icon = "👤" if turn["role"] == "user" else "🤖"
    print(f"{icon} {turn['message']}")

print("\n✅ Additional conversation added to session!")

print("\n🧠 Generating memories from extended conversation...\n")
print("⏳ Analyzing the full conversation history and consolidating memories...")

# Generate additional memories
operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    vertex_session_source={"session": session_name},
    config={"wait_for_completion": True},
)

print("\n✅ Memory generation complete!")

# Display newly generated memories
display_generated_memories(operation, client, "New Memories Generated")

print("COMPARISON: Scope-Based vs. Similarity-Based Retrieval")

# Method 1: Scope-based retrieval (returns ALL memories)
print("\n📊 Method 1: Scope-Based Retrieval")
print("Query: Get all memories for Emma")
print("-" * 80)

all_results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": guest_id}
)

all_memories = list(all_results)
print(f"\n✅ Found {len(all_memories)} total memories:\n")

for i, mem in enumerate(all_memories, 1):
    print(f"   {i}. {mem.memory.fact}")

# Method 2: Similarity search (returns only RELEVANT memories)
print("\n🔍 Method 2: Similarity Search")
print("Query: What are Emma's dietary restrictions?")
print("-" * 80)

search_results = client.agent_engines.memories.retrieve(
    name=agent_engine_name,
    scope={"user_id": guest_id},
    similarity_search_params={
        "search_query": "What are Emma's dietary restrictions?",
        "top_k": 3,  # Get top 3 most relevant memories
    },
)

relevant_memories = list(search_results)
print(f"\n✅ Found {len(relevant_memories)} relevant memories:\n")

for i, mem in enumerate(relevant_memories, 1):
    distance = mem.distance if hasattr(mem, "distance") else "N/A"
    print(f"   {i}. {mem.memory.fact}")
    print(f"      → Relevance distance: {distance} (lower = more relevant)")

print("\n🏨 Testing Similarity Search with Common Front Desk Questions\n")

# Test with various specific queries
search_queries = [
    "What are Emma's room location preferences?",
    "What are Emma's dietary restrictions?",
    "What spa services does Emma prefer?",
    "What temperature does Emma like in her room?",
]

for query in search_queries:
    print(f'\n❓ Query: "{query}"')
    print("-" * 80)

    # Perform similarity search
    results = client.agent_engines.memories.retrieve(
        name=agent_engine_name,
        scope={"user_id": guest_id},
        similarity_search_params={
            "search_query": query,
            "top_k": 2,  # Get top 2 most relevant memories
        },
    )

    memories = list(results)

    if memories:
        print("\n🎯 Most Relevant Memories:")
        for i, mem in enumerate(memories, 1):
            distance = mem.distance if hasattr(mem, "distance") else "N/A"
            print(f"   {i}. {mem.memory.fact}")
            print(
                f"      (Distance: {distance:.4f})"
                if isinstance(distance, (int, float))
                else f"      (Distance: {distance})"
            )
    else:
        print("   No relevant memories found")

    print()

print("✅ Similarity search enables fast, targeted information retrieval!")

print("🧹 Cleaning up resources...\n")

delete_agent_engine = True

if delete_agent_engine:
    # Delete the agent engine and all its memories
    client.agent_engines.delete(name=agent_engine_name, force=True)

    print("✅ Agent Engine deleted successfully!")
    print("   All guest memories have been removed")
else:
    print("⏭️  Skipping cleanup - Agent Engine will remain active")
    print(f"   Resource: {agent_engine_name}")
    print("\n⚠️  Remember to delete this resource later to avoid charges")