from __future__ import annotations

import pytest

from agent_framework.guardrails.framework_llm_client import classify_with_framework_llm
from agent_framework.guardrails.calibrated.llm_client import _BINARY_BLOCK_DIGIT, _BINARY_TASKS
from agent_framework.guardrails.calibrated.prompts.pinj import build_pinj_prompt
from agent_framework.guardrails.calibrated.prompts.revprec import build_revprec_prompt


class FakeLLM:
    def __init__(self, content: str):
        self.content = content
        self.profile_resolver = None

    async def ainvoke(self, messages, **kwargs):
        return self.content


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("task", "content", "allowed"),
    [
        ("PINJ", "1", True),
        ("PINJ", "0", False),
        ("PINJ", "prosa 1 e veredito 0", False),
        ("PINJ", "sem digito", True),
        ("PINJ", "", True),
        ("REVPREC", "1", False),
        ("REVPREC", "0", True),
        ("REVPREC", "Resposta: 1", False),
        ("REVPREC", "não sei", True),
        ("REVPREC", "", True),
    ],
)
async def test_binary_guardrail_semantics(task: str, content: str, allowed: bool, monkeypatch):
    monkeypatch.setenv("USE_MOCK_LLM", "false")
    out = await classify_with_framework_llm(
        FakeLLM(content), task, {"text": "texto", "context": {}}, profile_name="guardrail"
    )
    assert bool(out["allowed"]) is allowed


def test_binary_polarity_contract():
    assert _BINARY_BLOCK_DIGIT["REVPREC"] == "1"
    assert _BINARY_BLOCK_DIGIT.get("PINJ", "0") == "0"
    assert "PINJ" in _BINARY_TASKS and "REVPREC" in _BINARY_TASKS


def test_pinj_prompt_is_binary_and_keeps_examples():
    prompt = build_pinj_prompt("texto qualquer")
    assert "Responda APENAS um caractere" in prompt
    assert '"allowed"' not in prompt
    assert prompt.count("  Texto: ") == 17
    assert prompt.count("  Saída: 0") == 11
    assert prompt.count("  Saída: 1") == 6


def test_revprec_prompt_preserves_boundaries():
    prompt = build_revprec_prompt("FALA_NOVA", "\nHistorico da conversa:\n[assistant] CANCELEI_ANTES\n")
    assert "APENAS 1 ou 0" in prompt
    assert 'somente a fala do bloco "Resposta:"' in prompt
    assert "PROTOCOLO" in prompt
    assert "DESCRIÇÃO DA FATURA" in prompt