"""Tests for NobulexComplianceMiddleware."""

from langchain_nobulex.base import ModelResponse, ToolCall
from langchain_nobulex.middleware import NobulexComplianceMiddleware


class TestMiddlewareRuleEnforcement:
    def test_permit_allows_tool_call(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        response = ModelResponse(
            tool_calls=[ToolCall(name="read", args={"file": "data.csv"})]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].name == "read"

    def test_forbid_blocks_tool_call(self):
        mw = NobulexComplianceMiddleware(rules="forbid delete;")
        response = ModelResponse(
            tool_calls=[ToolCall(name="delete", args={"id": "123"})]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 0

    def test_forbid_with_condition_blocks_matching(self):
        mw = NobulexComplianceMiddleware(
            rules="forbid transfer where amount > 500;"
        )
        response = ModelResponse(
            tool_calls=[ToolCall(name="transfer", args={"amount": 1000})]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 0

    def test_forbid_with_condition_allows_non_matching(self):
        mw = NobulexComplianceMiddleware(
            rules="forbid transfer where amount > 500;"
        )
        response = ModelResponse(
            tool_calls=[ToolCall(name="transfer", args={"amount": 100})]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 1

    def test_default_allow_unknown_tool(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        response = ModelResponse(
            tool_calls=[ToolCall(name="unknown_tool", args={})]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 1

    def test_mixed_rules(self):
        mw = NobulexComplianceMiddleware(
            rules="permit read; forbid transfer where amount > 500; require log_all;"
        )
        response = ModelResponse(
            tool_calls=[
                ToolCall(name="read", args={"file": "a.txt"}),
                ToolCall(name="transfer", args={"amount": 1000}),
                ToolCall(name="transfer", args={"amount": 100}),
            ]
        )
        result = mw.after_model(response)
        assert len(result.tool_calls) == 2
        assert result.tool_calls[0].name == "read"
        assert result.tool_calls[1].name == "transfer"
        assert result.tool_calls[1].args["amount"] == 100


class TestMiddlewareAudit:
    def test_audit_records_all_actions(self):
        mw = NobulexComplianceMiddleware(
            rules="permit read; forbid delete;"
        )
        response = ModelResponse(
            tool_calls=[
                ToolCall(name="read", args={}),
                ToolCall(name="delete", args={"id": "1"}),
            ]
        )
        mw.after_model(response)
        log = mw.get_audit_log()
        assert len(log) == 2

    def test_audit_marks_allowed_blocked(self):
        mw = NobulexComplianceMiddleware(rules="forbid delete;")
        response = ModelResponse(
            tool_calls=[
                ToolCall(name="read", args={}),
                ToolCall(name="delete", args={}),
            ]
        )
        mw.after_model(response)
        log = mw.get_audit_log()
        read_entry = next(e for e in log if e["action"] == "read")
        delete_entry = next(e for e in log if e["action"] == "delete")
        assert read_entry["allowed"] is True
        assert delete_entry["allowed"] is False

    def test_audit_records_matching_rule(self):
        mw = NobulexComplianceMiddleware(
            rules="forbid transfer where amount > 500;"
        )
        response = ModelResponse(
            tool_calls=[ToolCall(name="transfer", args={"amount": 1000})]
        )
        mw.after_model(response)
        log = mw.get_audit_log()
        assert "forbid transfer" in log[0]["matching_rule"]

    def test_audit_hash_chain(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        response = ModelResponse(
            tool_calls=[
                ToolCall(name="read", args={"f": "1"}),
                ToolCall(name="read", args={"f": "2"}),
                ToolCall(name="read", args={"f": "3"}),
            ]
        )
        mw.after_model(response)
        log = mw.get_audit_log()
        assert log[1]["prev_hash"] == log[0]["hash"]
        assert log[2]["prev_hash"] == log[1]["hash"]

    def test_get_audit_log_json(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        response = ModelResponse(
            tool_calls=[ToolCall(name="read", args={})]
        )
        mw.after_model(response)
        json_str = mw.get_audit_log_json()
        assert '"action": "read"' in json_str


class TestMiddlewareAfterAgent:
    def test_after_agent_attaches_log(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        response = ModelResponse(
            tool_calls=[ToolCall(name="read", args={})]
        )
        mw.after_model(response)
        state = mw.after_agent({})
        assert "audit_log" in state
        assert len(state["audit_log"]) == 1

    def test_after_agent_preserves_existing_state(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        state = mw.after_agent({"existing_key": "value"})
        assert state["existing_key"] == "value"
        assert "audit_log" in state


class TestMiddlewareWrapToolCall:
    def test_wrap_tool_call_passes_through(self):
        mw = NobulexComplianceMiddleware(rules="permit read;")
        tc = ToolCall(name="read", args={})
        result = mw.wrap_tool_call(tc)
        assert result is tc
