"""Unit tests for coalex.ext.guardrail."""

from __future__ import annotations

import pytest
from coalex.ext.guardrail import guardrail_span
from opentelemetry.trace import StatusCode


def test_guardrail_span_sets_kind(memory_exporter):
    @guardrail_span()
    def check_toxicity(text: str) -> bool:
        return True

    check_toxicity(text="hello")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "GUARDRAIL"


def test_guardrail_span_sets_name_from_function(memory_exporter):
    @guardrail_span()
    def pii_filter(text: str) -> bool:
        return True

    pii_filter(text="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.name"] == "pii_filter"


def test_guardrail_span_sets_custom_guardrail_name(memory_exporter):
    @guardrail_span(name="pii_check", guardrail_name="PII Filter v2")
    def check(text: str) -> bool:
        return True

    check(text="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].name == "pii_check"
    assert spans[0].attributes["guardrail.name"] == "PII Filter v2"


def test_guardrail_span_result_pass_on_success(memory_exporter):
    @guardrail_span()
    def safe_check(text: str) -> bool:
        return True

    safe_check(text="safe content")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.result"] == "pass"


def test_guardrail_span_result_fail_on_exception(memory_exporter):
    @guardrail_span()
    def strict_check(text: str) -> None:
        raise ValueError("content blocked")

    with pytest.raises(ValueError, match="content blocked"):
        strict_check(text="bad content")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.result"] == "fail"


def test_guardrail_span_exception_is_reraised(memory_exporter):
    @guardrail_span()
    def check(text: str) -> None:
        raise RuntimeError("blocked")

    with pytest.raises(RuntimeError, match="blocked"):
        check(text="x")


def test_guardrail_span_score_set_when_return_is_float(memory_exporter):
    @guardrail_span()
    def toxicity_score(text: str) -> float:
        return 0.12

    result = toxicity_score(text="hello world")
    assert result == 0.12
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.score"] == 0.12
    assert spans[0].attributes["guardrail.result"] == "pass"


def test_guardrail_span_score_not_set_for_non_float_return(memory_exporter):
    @guardrail_span()
    def check(text: str) -> bool:
        return True

    check(text="ok")
    spans = memory_exporter.get_finished_spans()
    assert "guardrail.score" not in spans[0].attributes


def test_guardrail_span_returns_function_result(memory_exporter):
    @guardrail_span()
    def check(text: str) -> str:
        return "approved"

    result = check(text="safe")
    assert result == "approved"


@pytest.mark.asyncio
async def test_guardrail_span_async_sets_kind(memory_exporter):
    @guardrail_span(name="async_toxicity")
    async def check(text: str) -> bool:
        return True

    await check(text="hello")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "GUARDRAIL"
    assert spans[0].attributes["guardrail.result"] == "pass"


@pytest.mark.asyncio
async def test_guardrail_span_async_result_fail_on_exception(memory_exporter):
    @guardrail_span()
    async def strict(text: str) -> None:
        raise ValueError("async blocked")

    with pytest.raises(ValueError, match="async blocked"):
        await strict(text="bad")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.result"] == "fail"


@pytest.mark.asyncio
async def test_guardrail_span_async_score(memory_exporter):
    @guardrail_span()
    async def score_fn(text: str) -> float:
        return 0.05

    result = await score_fn(text="safe text")
    assert result == 0.05
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["guardrail.score"] == 0.05


# ---------------------------------------------------------------------------
# Status OK / ERROR tests
# ---------------------------------------------------------------------------


def test_guardrail_span_status_ok_on_success(memory_exporter):
    @guardrail_span()
    def check(text: str) -> bool:
        return True

    check(text="safe")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


def test_guardrail_span_status_error_on_exception(memory_exporter):
    @guardrail_span()
    def check(text: str) -> None:
        raise ValueError("blocked")

    with pytest.raises(ValueError, match="blocked"):
        check(text="bad")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    assert "blocked" in spans[0].status.description
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1


@pytest.mark.asyncio
async def test_guardrail_span_async_status_ok(memory_exporter):
    @guardrail_span()
    async def check(text: str) -> bool:
        return True

    await check(text="safe")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


@pytest.mark.asyncio
async def test_guardrail_span_async_status_error(memory_exporter):
    @guardrail_span()
    async def check(text: str) -> None:
        raise ValueError("async blocked")

    with pytest.raises(ValueError, match="async blocked"):
        await check(text="bad")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1
