"""NobulexComplianceMiddleware — covenant-based compliance for LangChain agents."""

from __future__ import annotations

from typing import Any

from langchain_nobulex.audit import AuditLog
from langchain_nobulex.base import AgentMiddleware, ModelResponse, ToolCall
from langchain_nobulex.parser import Rule, parse_rules


class NobulexComplianceMiddleware(AgentMiddleware):
    """Compliance middleware that enforces covenant rules on agent tool calls.

    Parses permit/forbid/require rules, intercepts tool calls, blocks
    forbidden actions, and maintains a hash-chained audit trail.

    Args:
        rules: A covenant rules string, e.g.
            "permit read; forbid transfer where amount > 500; require log_all;"
    """

    def __init__(self, rules: str) -> None:
        self.rules: list[Rule] = parse_rules(rules)
        self.audit: AuditLog = AuditLog()

    def _evaluate_tool_call(self, tool_call: ToolCall) -> tuple[bool, str]:
        """Check a tool call against all rules.

        Returns:
            (allowed, matching_rule_str)
        """
        for rule in self.rules:
            if rule.action == "forbid" and rule.matches(tool_call.name, tool_call.args):
                return False, str(rule)

        for rule in self.rules:
            if rule.action == "permit" and rule.matches(tool_call.name, tool_call.args):
                return True, str(rule)

        # Default: allow if no forbid matched
        return True, "default:allow"

    def after_model(self, response: ModelResponse) -> ModelResponse:
        """Intercept tool calls, check against rules, block forbidden ones."""
        allowed_calls: list[ToolCall] = []

        for tc in response.tool_calls:
            is_allowed, matching_rule = self._evaluate_tool_call(tc)

            self.audit.add_entry(
                action=tc.name,
                params=tc.args,
                allowed=is_allowed,
                matching_rule=matching_rule,
            )

            if is_allowed:
                allowed_calls.append(tc)

        response.tool_calls = allowed_calls
        return response

    def wrap_tool_call(self, tool_call: ToolCall) -> ToolCall | None:
        """Log tool execution. Called for each tool that passed after_model."""
        # Already logged in after_model; this hook is for additional
        # per-execution logging if needed.
        return tool_call

    def after_agent(self, state: dict[str, Any]) -> dict[str, Any]:
        """Attach the complete audit log to the agent's final state."""
        state["audit_log"] = self.audit.export_dicts()
        return state

    def get_audit_log(self) -> list[dict[str, Any]]:
        """Export the full tamper-proof audit chain."""
        return self.audit.export_dicts()

    def get_audit_log_json(self) -> str:
        """Export the full audit chain as formatted JSON."""
        return self.audit.export_json()
