import pytest

from agent_framework.guardrails.base import RailDecision
from agent_framework.guardrails.output_supervisor import OutputSupervisor
from agent_framework.guardrails.parallel_executor import ParallelRailExecutor
from agent_framework.guardrails.rail_action import RailAction
from agent_framework.observability.code_mapper import ObservabilityCodeMapper


def test_compact_and_rich_mapping_are_backward_compatible():
    mapper = ObservabilityCodeMapper({
        "guardrail.dlex_in": "GRL.004",
        "guardrail.tox": {"label": "GRL.005", "aliases": ["TOX"]},
    })
    assert mapper.map("guardrail.dlex_in") == "GRL.004"
    assert mapper.map("DLEX_IN") == "GRL.004"
    assert mapper.map("TOX") == "GRL.005"
    assert mapper.map("guardrail.unknown") == "guardrail.unknown"


def test_action_resolution_accepts_internal_and_external_aliases():
    mapper = ObservabilityCodeMapper({
        "guardrail.revprec": {
            "action": "retry",
            "aliases": ["REVPREC", "TIM_REVPREC"],
        },
        "guardrail.handover": {
            "action": "handover",
            "aliases": ["ATH", "HUMAN"],
        },
    })
    assert mapper.action_for("REVPREC") == "retry"
    assert mapper.action_for("TIM_REVPREC") == "retry"
    assert mapper.action_for("guardrail.revprec") == "retry"
    assert mapper.action_for("ATH") == "handover"
    assert mapper.action_for("UNKNOWN") is None


class _DeniedLegacyRail:
    code = "REVPREC"

    async def evaluate(self, text, context):
        return RailDecision(code=self.code, allowed=False, reason="premature")


class _DeniedPolicyRail:
    code = "REVPREC"
    _guardrail_policy = {"on_deny": "handover"}

    async def evaluate(self, text, context):
        return RailDecision(code=self.code, allowed=False, reason="premature")


@pytest.mark.asyncio
async def test_output_supervisor_uses_registry_action_without_name_hardcode():
    mapper = ObservabilityCodeMapper({"guardrail.revprec": {"action": "retry"}})
    supervisor = OutputSupervisor(
        rails=[_DeniedLegacyRail()],
        enable_parallel=False,
        observability_mapper=mapper,
    )
    decision = await supervisor.evaluate("candidate", {})
    assert decision.action == RailAction.RETRY
    assert decision.results[0].metadata.get("action_source") == "observability_mapping"


@pytest.mark.asyncio
async def test_guardrails_yaml_policy_precedes_registry_action():
    mapper = ObservabilityCodeMapper({"guardrail.revprec": {"action": "retry"}})
    supervisor = OutputSupervisor(
        rails=[_DeniedPolicyRail()],
        enable_parallel=False,
        observability_mapper=mapper,
    )
    decision = await supervisor.evaluate("candidate", {})
    assert decision.action == RailAction.HANDOVER


@pytest.mark.asyncio
async def test_parallel_executor_uses_same_registry():
    mapper = ObservabilityCodeMapper({"guardrail.revprec": {"action": "retry"}})
    executor = ParallelRailExecutor(
        fail_fast=True,
        observability_mapper=mapper,
    )
    execution = await executor.run("candidate", {}, [_DeniedLegacyRail()])
    assert execution.terminal_result is not None
    assert execution.terminal_result.action == RailAction.RETRY
    assert execution.terminal_result.metadata.get("action_source") == "observability_mapping"