"""Tests for the main Shield class."""

import os
import sys
sys.path.insert(0, ".")

from agent_shield import Shield

AUDIT = "test_shield_audit.jsonl"
CONSENT = "test_shield_consent.jsonl"


def setup():
    for f in [AUDIT, CONSENT]:
        if os.path.exists(f):
            os.remove(f)


def teardown():
    for f in [AUDIT, CONSENT]:
        if os.path.exists(f):
            os.remove(f)


def test_scan():
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    result = shield.scan("Email me at test@example.com")
    assert result.has_pii
    assert "EMAIL" in result.pii_found
    assert "[EMAIL_REDACTED]" in result.redacted
    assert result.original == "Email me at test@example.com"
    teardown()


def test_scan_clean():
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    result = shield.scan("No personal data here")
    assert not result.has_pii
    assert result.pii_found == {}
    teardown()


def test_consent_lifecycle():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    shield.consent.grant("u1", "marketing", method="opt_in")
    assert shield.consent.check("u1", "marketing") is True

    shield.consent.withdraw("u1", "marketing")
    assert shield.consent.check("u1", "marketing") is False

    assert shield.consent.check("u1", "nonexistent") is False
    teardown()


def test_generic_wrapper():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)

    def echo(text):
        return f"Echo: {text}"

    result = shield.call(echo, "My email is x@y.com", provider="test", purpose="demo")
    assert "Echo:" in str(result.response)
    assert "EMAIL" in result.pii_detected

    valid, count = shield.verify_audit()
    assert valid
    assert count == 1
    teardown()


def test_dpia_generation():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    shield.audit.record(provider="openai", model="gpt-4", purpose="support",
                        input_text="test@test.com", output_text="ok",
                        tokens_in=5, tokens_out=2)

    dpia = shield.generate_dpia(system_name="Test")
    assert "Data Protection Impact Assessment" in dpia
    assert "Test" in dpia
    assert "openai" in dpia.lower()
    assert "januscompliance.co.uk" in dpia
    teardown()


def test_dataflow_generation():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    shield.audit.record(provider="anthropic", model="claude", purpose="analysis",
                        input_text="data", output_text="result")

    flow = shield.generate_dataflow()
    assert "anthropic" in flow.lower()
    assert "United States" in flow
    teardown()


def test_mermaid_generation():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    shield.audit.record(provider="openai", input_text="hi", output_text="hello")

    diagram = shield.generate_dataflow_diagram("MyApp")
    assert "graph LR" in diagram
    assert "MyApp" in diagram
    teardown()


def test_user_report():
    setup()
    shield = Shield(audit_path=AUDIT, consent_path=CONSENT)
    shield.consent.grant("u1", "support")
    shield.audit.record(provider="test", user_id="u1",
                        input_text="query", output_text="response")

    report = shield.user_report("u1")
    assert report["user_id"] == "u1"
    assert len(report["audit_entries"]) == 1
    assert len(report["consent_history"]) == 1
    assert "support" in report["active_consents"]
    teardown()


if __name__ == "__main__":
    for name, fn in list(globals().items()):
        if name.startswith("test_") and callable(fn):
            try:
                fn()
                print(f"  PASS {name}")
            except AssertionError as e:
                print(f"  FAIL {name}: {e}")
    print("Done.")
