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

%pip install "google-cloud-aiplatform[agent_engines]" --force-reinstall --quiet
# Install the Google Generative AI integration for LangChain
%pip install -U "langchain-google-genai" --quiet

# Restart the session to ensure packages are updated
import os

os._exit(0)

import sys

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

    auth.authenticate_user()

import os

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
LOCATION = "us-central1" # @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"))
if not LOCATION:
    LOCATION = os.environ.get("GOOGLE_CLOUD_REGION")
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"

import datetime

import requests
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import (
    BaseMessage,
    HumanMessage,
    message_to_dict,
    messages_from_dict,
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from vertexai import Client

class VertexAISessionChatMessageHistory(BaseChatMessageHistory):
    """A LangChain message history backend that defers storage to
    Vertex AI Agent Engine's Session Service via the vertexai Client.
    """

    def __init__(self, client: Client, session_name: str, author: str = "user"):
        self.client = client
        self.session_name = session_name
        self.author = author

    @property
    def messages(self) -> list[BaseMessage]:
        # Call ListEvents from the Vertex AI Session service to return the list of messages.
        events = self.client.agent_engines.sessions.events.list(name=self.session_name)

        lc_messages = []
        for event in events:
            # Convert directly from the LangChain message stored in raw_event
            if event.raw_event and "langchain_message" in event.raw_event:
                # `messages_from_dict` expects a list of serialized messages
                extracted_message = messages_from_dict(
                    [event.raw_event["langchain_message"]]
                )[0]
                lc_messages.append(extracted_message)
                continue
        return lc_messages

    def add_message(self, message: BaseMessage) -> None:
        is_human = isinstance(message, HumanMessage)
        author = self.author if is_human else "agent"

        # Dump the entire LangChain message structure (id, tool_calls, kwargs, response_metadata)
        # into a JSON serializable dict
        serialized_lc_message = message_to_dict(message)

        self.client.agent_engines.sessions.events.append(
            name=self.session_name,
            author=author,
            invocation_id="langchain-invocation",
            timestamp=datetime.datetime.now(datetime.timezone.utc),
            config={
                # Store the full LangChain payload in the raw_event field
                "raw_event": {"langchain_message": serialized_lc_message}
            },
        )

    def clear(self) -> None:
        """Clears the session. Vertex AI Session Service currently doesn't support
        easily clearing individual events; typical workarounds include
        deleting the session natively via `self.client.agent_engines.sessions.delete(...)`.
        """
        raise NotImplementedError(
            "Clear is not natively supported without deleting the entire session via SDK."
        )

import vertexai

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

# Create an agent engine to hold the session
agent_engine = client.agent_engines.create(config={"display_name": "langchain_test"})
agent_engine_path = agent_engine.api_resource.name

# Create a Vertex AI session resource
session_operation = client.agent_engines.sessions.create(
    name=agent_engine_path, user_id="langchain_user"
)

# Extract the full session resource name from the finished operation response
session_resource_name = session_operation.response.name

# Define a factory function that LangChain will call per session ID
def get_vertex_session_history(session_id: str) -> BaseChatMessageHistory:
    return VertexAISessionChatMessageHistory(client=client, session_name=session_id)

# Create a standard LangChain Model and Prompt
model = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash", project=PROJECT_ID, location=LOCATION
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant."),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

chain = prompt | model

# Wrap the chain with RunnableWithMessageHistory
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_vertex_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Invoke the chain using the session resource name we generated
input = "Hi! My name is John Doe."
print(input)
response = chain_with_history.invoke(
    {"input": input}, config={"configurable": {"session_id": session_resource_name}}
)
print("Gemini: ", response.content)

input = "What is my name?"
print(input)
response = chain_with_history.invoke(
    {"input": input}, config={"configurable": {"session_id": session_resource_name}}
)
print("Gemini: ", response.content)

# View the conversation history
history = get_vertex_session_history(session_resource_name)
for message in history.messages:
    print(f"{message.type.capitalize()}: {message.content}")

# Create a new Vertex AI session resource
session_operation = client.agent_engines.sessions.create(
    name=agent_engine_path, user_id="langchain_user"
)

# Extract the full session resource name from the finished operation response
session_resource_name = session_operation.response.name

# Define the tool
@tool
def get_weather(location: str) -> str:
    """Returns the current weather for a given city name."""
    # format=3 gives a simple string like: "London: ⛅️ +15°C"
    try:
        response = requests.get(f"https://wttr.in/{location}?format=3")
        return response.text.strip()
    except Exception:
        return "Sorry, I couldn't fetch the weather right now."

# Bind the tool to the model
# ChatVertexAI natively supports LangChain tool integration
model_with_tools = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash", project=PROJECT_ID, location=LOCATION
).bind_tools([get_weather])

# Create a Prompt Template with STRICT system instructions
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a specialized weather assistant. You MUST use the `get_weather` tool to answer any questions about the weather. Do not use your internal knowledge.",
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

# Combine the prompt and the model into a standard runnable chain
chain = prompt | model_with_tools

# Wrap the chain with RunnableWithMessageHistory
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_vertex_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# --- Turn 1: AI Uses the Tool ---
input = "What is the weather in Paris?"
print(input)
response = chain_with_history.invoke(
    {"input": input}, config={"configurable": {"session_id": session_resource_name}}
)

# If the system instruction worked, the model will output a tool call instead of text
print("(Gemini Tool Request):", response.tool_calls)

# --- Turn 2: Providing the tool result ---
if response.tool_calls:
    tool_call = response.tool_calls[0]
    # Execute the python function manually
    tool_result = get_weather.invoke(tool_call)
    print("(Tool Response):", tool_result)
    request_payload = {
        # Pass the ToolMessage containing the tool_call_id directly into the input
        "input": [tool_result]
    }
    final_response = chain_with_history.invoke(
        request_payload, config={"configurable": {"session_id": session_resource_name}}
    )
    print("(Gemini Final Answer):", final_response.content)

history = get_vertex_session_history(session_resource_name)
for message in history.messages:
    print(f"{message.type.capitalize()}: {message.content}")

client.agent_engines.delete(name=agent_engine_path, force=True)