# 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 "crewai" "google-cloud-aiplatform>=1.100.0"

# import sys

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

#     auth.authenticate_user()

import json
import os
import uuid

import vertexai

# Set project configuration
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
    if not PROJECT_ID:
        raise ValueError("Project ID not found. Please set it in the form above or as the GOOGLE_CLOUD_PROJECT environment variable.")

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
SERVICE_ACCOUNT_FILE = "[your-service-account-file]"  # @param {type:"string"}

# Set environment variables for CrewAI
os.environ["DEFAULT_VERTEXAI_PROJECT"] = PROJECT_ID
os.environ["DEFAULT_VERTEXAI_LOCATION"] = LOCATION

# Agent configuration
MODEL_NAME = "gemini-2.5-flash"
USER_ID = f"user_{uuid.uuid4()}"

print(f"Project: {PROJECT_ID}")
print(f"Location: {LOCATION}")
print(f"Session User ID: {USER_ID}")

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

# Load the service account credentials with Vertex AI User permission for the CrewAI LLM
with open(SERVICE_ACCOUNT_FILE, "r") as file:
    vertex_credentials = json.load(file)
vertex_credentials_json = json.dumps(vertex_credentials)

from datetime import datetime
from typing import Any, Dict, List
from crewai import LLM, Agent, Crew, Task
from crewai.memory.external.external_memory import ExternalMemory
from crewai.memory.storage.interface import Storage

def run_single_turn(query: str, agent: Agent, memory: ExternalMemory):
    """Runs a single conversational turn with the CrewAI agent."""
    task = Task(
        description=f"Respond to the user's query: '{query}'",
        expected_output="A helpful and personalized response based on your memory of the user.",
        agent=agent,
    )
    crew = Crew(
        agents=[agent],
        tasks=[task],
        external_memory=memory,
    )
    return crew.kickoff()


def chat_loop(agent: Agent, memory: ExternalMemory) -> None:
    """Main chat interface loop."""
    print("\nStarting chat. Type 'exit' or 'quit' to end.")
    print("Every message will be automatically stored in memory.\n")

    while True:
        user_input = input("\nYou: ")
        if user_input.lower() in ["quit", "exit", "bye"]:
            print("\nAssistant: Thank you for chatting. Have a great day!")
            break

        response = run_single_turn(user_input, agent, memory)
        print(f"\nAssistant: {response}")

agent_engine = client.agent_engines.create()
print(f"Created Agent Engine: {agent_engine.api_resource.name}")

import_initial_memory_op = client.agent_engines.create_memory(
    name=agent_engine.api_resource.name,
    fact=f"The user's name is {USER_ID} and they have a passion for learning new things.",
    scope={"user_id": USER_ID},
)

# Pretty print the created memory
memory = import_initial_memory_op.response
print("Created initial memory:")
print(f"  Fact: {memory.fact}")
print(f"  User ID: {memory.scope.get('user_id')}")
print(f"  Created at: {memory.create_time}")

memories = client.agent_engines.retrieve_memories(
    name=agent_engine.api_resource.name,
    scope={"user_id": USER_ID},
)

for i, memory in enumerate(memories, 1):
    if hasattr(memory, "memory") and hasattr(memory.memory, "fact"):
        print(f"{i}. {memory.memory.fact}")

class VertexAIMemoryBankStorage(Storage):
    """CrewAI storage implementation using Vertex AI Memory Bank."""

    def __init__(
        self,
        project: str,
        location: str,
        agent_engine_name: str,
        use_semantic_search: bool = True,
        default_user_id: str = "default_user",
    ):
        """Initialize Vertex AI Memory Bank storage."""
        self.project = project
        self.location = location
        self.agent_engine_name = agent_engine_name
        self.use_semantic_search = use_semantic_search
        self.default_user_id = default_user_id

        # Validate required parameters early
        if not all([project, location, agent_engine_name]):
            raise ValueError("project, location, and agent_engine_name are required")

    def save(self, value: str, metadata: Dict[str, Any]) -> None:
        """Save a memory to Vertex AI Memory Bank."""

        client = self._get_client()

        # Extract user_id from metadata or use default
        user_id = metadata.get("user_id", self.default_user_id)

        # Prepare memory event
        event = {
            "content": {
                "parts": [{"text": value}],
                "role": metadata.get("role", "user"),
            }
        }

        # Add timestamp if provided
        if "timestamp" in metadata:
            event["timestamp"] = metadata["timestamp"]

        # Generate memory in Vertex AI
        client.agent_engines.generate_memories(
            name=self.agent_engine_name,
            direct_contents_source={"events": [event]},
            scope={"user_id": user_id},
            config={"wait_for_completion": True},
        )

    def search(self, query: str, limit: int = 5, **kwargs) -> List[Dict[str, Any]]:
        """Search memories in Vertex AI Memory Bank.

        Args:
            query: Search query (used for semantic search if enabled)
            limit: Maximum number of results to return
            **kwargs: Additional parameters (e.g., user_id)

        Returns:
            List of memory dictionaries with 'value' and 'metadata' keys
        """
        client = self._get_client()

        # Extract user_id from kwargs or use default
        user_id = kwargs.get("user_id", self.default_user_id)
        scope = {"user_id": user_id}

        # Determine search method
        if self.use_semantic_search and query:
            retrieved_memories = client.agent_engines.retrieve_memories(
                name=self.agent_engine_name,
                scope=scope,
                similarity_search_params={"search_query": query, "top_k": limit},
            )
        else:
            retrieved_memories = client.agent_engines.retrieve_memories(
                name=self.agent_engine_name, scope=scope, simple_retrieval_params={}
            )

        # Transform results to CrewAI format
        results = []
        for memory in retrieved_memories:
            if hasattr(memory, "memory") and hasattr(memory.memory, "fact"):
                result = {
                    "memory": memory.memory.fact,
                    "metadata": {
                        "user_id": user_id,
                        "timestamp": (
                            memory.memory.update_time.isoformat()
                            if hasattr(memory.memory, "update_time")
                            else datetime.now().isoformat()
                        ),
                    },
                }

                results.append(result)

        return results

    def reset(self):
        client = self._get_client()
        client.agent_engines.delete(name=self.agent_engine_name)

    def _get_client(self):
        """Create a Vertex AI client instance.

        Returns a new client for each operation to ensure proper
        event loop management in async contexts.
        """
        return vertexai.Client(project=self.project, location=self.location)

# Instantiate the LLM
llm = LLM(model=MODEL_NAME, temperature=0.7, vertex_credentials=vertex_credentials_json)

# Instantiate your custom storage and the external memory
vertex_memory_bank_storage = VertexAIMemoryBankStorage(
    project=PROJECT_ID,
    location=LOCATION,
    agent_engine_name=agent_engine.api_resource.name,
    default_user_id=USER_ID,
)
external_vertex_memory_bank = ExternalMemory(storage=vertex_memory_bank_storage)

# Create a personal assistant agent
assistant = Agent(
    role="Personal Assistant",
    goal="Remember user information and provide personalized help",
    backstory="You are a helpful assistant with perfect memory of past conversations.",
    llm=llm,
)

chat_loop(assistant, external_vertex_memory_bank)

memories = client.agent_engines.retrieve_memories(
    name=agent_engine.api_resource.name,
    scope={"user_id": USER_ID},
)

for i, memory in enumerate(memories, 1):
    if hasattr(memory, "memory") and hasattr(memory.memory, "fact"):
        print(f"{i}. {memory.memory.fact}")

print(f"\nTotal memories stored: {len(list(memories))}")

delete_engine = True

if delete_engine:
    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)