from typing import Any

from langchain.agents import create_agent
from langchain.agents.middleware import PIIMiddleware
from langchain_core.language_models.chat_models import SimpleChatModel
from langchain_core.messages import BaseMessage


class EchoLastUserMessage(SimpleChatModel):
    @property
    def _llm_type(self) -> str:
        return "echo-last-user-message"

    def bind_tools(self, tools: list[Any], *, tool_choice: Any = None, **kwargs: Any) -> "EchoLastUserMessage":
        if tools:
            raise NotImplementedError("This validation model does not exercise tool calls.")
        return self

    def _call(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> str:
        user_messages = [message for message in messages if message.type == "human"]
        content = user_messages[-1].content if user_messages else ""
        if not isinstance(content, str):
            content = str(content)
        return f"model saw: {content}"


agent = create_agent(
    model=EchoLastUserMessage(),
    tools=[],
    middleware=[
        PIIMiddleware("email", strategy="redact", apply_to_input=True),
    ],
)

prompt = "Send the receipt to priya@example.com."
result = agent.invoke({"messages": [{"role": "user", "content": prompt}]})
reply = result["messages"][-1].content

print("Input prompt:")
print(prompt)
print()
print("Agent reply:")
print(reply)
print()

if "priya@example.com" in reply:
    raise SystemExit("PII leaked to the model reply")
if "[REDACTED_EMAIL]" not in reply:
    raise SystemExit("Redaction marker was not present")

print("PII leaked to model: no")