from __future__ import annotations

import asyncio
from unittest.mock import MagicMock

from langchain_core.messages import ToolMessage

from deerflow.agents.middlewares.sensitive_output_redaction_middleware import SensitiveOutputRedactionMiddleware


def _request(name: str = "bash", args: dict | None = None):
    req = MagicMock()
    req.tool_call = {"name": name, "args": args or {}, "id": "call_1"}
    return req


def test_redacts_key_value_and_bearer():
    mw = SensitiveOutputRedactionMiddleware()
    req = _request()
    handler = MagicMock(
        return_value=ToolMessage(
            content="OPENAI_API_KEY=sk-abc123456789\nAuthorization: Bearer abcdefghijklmnop",
            tool_call_id="call_1",
            name="bash",
            status="success",
        )
    )
    result = mw.wrap_tool_call(req, handler)
    assert "OPENAI_API_KEY=[REDACTED]" in result.content
    assert "Bearer [REDACTED]" in result.content


def test_redacts_private_key_block():
    mw = SensitiveOutputRedactionMiddleware()
    req = _request()
    handler = MagicMock(
        return_value=ToolMessage(
            content="-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----",
            tool_call_id="call_1",
            name="bash",
            status="success",
        )
    )
    result = mw.wrap_tool_call(req, handler)
    assert result.content == "[REDACTED]"


def test_async_path_redacts_jwt():
    mw = SensitiveOutputRedactionMiddleware()
    req = _request()

    async def handler(_):
        return ToolMessage(
            content="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.abcde12345.zyxwv98765",
            tool_call_id="call_1",
            name="bash",
            status="success",
        )

    result = asyncio.run(mw.awrap_tool_call(req, handler))
    assert result.content == "[REDACTED]"