"""Redact sensitive values from tool outputs before they re-enter the model context."""

from __future__ import annotations

import re
from collections.abc import Awaitable, Callable
from typing import override

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import ToolMessage
from langgraph.prebuilt.tool_node import ToolCallRequest
from langgraph.types import Command

_SIMPLE_REDACTION = "[REDACTED]"

_PATTERNS = [
    re.compile(r"(?im)\b([A-Z][A-Z0-9_]{1,64})\s*=\s*([^\s\"'`]{6,})"),
    re.compile(r"(?i)\b(bearer\s+)[A-Za-z0-9._\-+/=]{8,}"),
    re.compile(r"(?i)\b(sk-[A-Za-z0-9]{12,})\b"),
    re.compile(r"(?i)\b(eyJ[A-Za-z0-9_\-]{10,}\.[A-Za-z0-9_\-]{10,}\.[A-Za-z0-9_\-]{10,})\b"),
    re.compile(r"(?is)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----"),
]


def _redact_text(text: str) -> str:
    value = text
    # Preserve var name for KEY=VALUE style output.
    value = _PATTERNS[0].sub(lambda m: f"{m.group(1)}={_SIMPLE_REDACTION}", value)
    value = _PATTERNS[1].sub(lambda m: f"{m.group(1)}{_SIMPLE_REDACTION}", value)
    for pattern in _PATTERNS[2:]:
        value = pattern.sub(_SIMPLE_REDACTION, value)
    return value


class SensitiveOutputRedactionMiddleware(AgentMiddleware[AgentState]):
    """Redact secrets from tool outputs."""

    def _redact_tool_result(self, result: ToolMessage | Command) -> ToolMessage | Command:
        if not isinstance(result, ToolMessage):
            return result
        if isinstance(result.content, str):
            redacted = _redact_text(result.content)
            if redacted != result.content:
                return ToolMessage(content=redacted, tool_call_id=result.tool_call_id, name=result.name, status=result.status)
            return result
        return result

    @override
    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], ToolMessage | Command],
    ) -> ToolMessage | Command:
        return self._redact_tool_result(handler(request))

    @override
    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
    ) -> ToolMessage | Command:
        return self._redact_tool_result(await handler(request))

