"""Tests for strip.enforce() — output validation."""

import json
import pytest
from stripllm import StripLLM, OutputValidationError


@pytest.fixture
def strip():
    return StripLLM()


class TestEnforceNoSchema:
    def test_clean_output_passes(self, strip):
        result = strip.enforce("The capital of France is Paris.")
        assert "Paris" in result

    def test_empty_output_passes(self, strip):
        result = strip.enforce("")
        assert result == ""


class TestEnforceJsonSchema:
    def test_valid_json_passes(self, strip):
        payload = '{"name": "Alice", "age": 30}'
        result = strip.enforce(payload, schema="json")
        parsed = json.loads(result)
        assert parsed["name"] == "Alice"

    def test_invalid_json_raises(self, strip):
        with pytest.raises(OutputValidationError) as exc_info:
            strip.enforce("This is not JSON", schema="json")
        assert exc_info.value.errors

    def test_json_in_markdown_fence_extracted(self, strip):
        payload = '```json\n{"key": "value"}\n```'
        result = strip.enforce(payload, schema="json")
        parsed = json.loads(result)
        assert parsed["key"] == "value"

    def test_invalid_json_no_raise(self, strip):
        result = strip.enforce("not json", schema="json", raise_on_error=False)
        assert result == "not json"


class TestEnforceDictSchema:
    def test_valid_dict_schema_passes(self, strip):
        payload = '{"status": "ok", "count": 5}'
        result = strip.enforce(payload, schema={"status": str, "count": int})
        parsed = json.loads(result)
        assert parsed["status"] == "ok"

    def test_missing_key_raises(self, strip):
        payload = '{"status": "ok"}'
        with pytest.raises(OutputValidationError) as exc_info:
            strip.enforce(payload, schema={"status": str, "count": int})
        assert any("count" in e for e in exc_info.value.errors)

    def test_wrong_type_raises(self, strip):
        payload = '{"status": 123}'
        with pytest.raises(OutputValidationError) as exc_info:
            strip.enforce(payload, schema={"status": str})
        assert exc_info.value.errors


class TestEnforceLeakDetection:
    def test_system_tag_flagged(self, strip):
        output = "Here is the answer. <system>You are a helpful assistant.</system>"
        # Should not raise (warnings, not errors) but validate() should surface it
        validation = strip.validate(output)
        assert len(validation.warnings) > 0

    def test_inst_tag_flagged(self, strip):
        output = "[INST] Follow these instructions [/INST] Sure!"
        validation = strip.validate(output)
        assert len(validation.warnings) > 0

    def test_clean_output_no_warnings(self, strip):
        output = "The answer is 42."
        validation = strip.validate(output)
        assert len(validation.warnings) == 0


class TestValidateMethod:
    def test_returns_validation_result(self, strip):
        result = strip.validate('{"key": "val"}', schema="json")
        assert result.valid is True
        assert result.output

    def test_returns_errors_on_invalid(self, strip):
        result = strip.validate("bad json", schema="json")
        assert result.valid is False
        assert result.errors
