"""Tests for audit logging."""

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

from agent_shield.audit import AuditLog

AUDIT_FILE = "test_audit_temp.jsonl"


def setup():
    if os.path.exists(AUDIT_FILE):
        os.remove(AUDIT_FILE)


def teardown():
    if os.path.exists(AUDIT_FILE):
        os.remove(AUDIT_FILE)


def test_record_and_query():
    setup()
    log = AuditLog(AUDIT_FILE)
    log.record(provider="openai", model="gpt-4", purpose="test",
               user_id="u1", input_text="hello", output_text="hi")
    log.record(provider="anthropic", model="claude", purpose="test",
               user_id="u2", input_text="bye", output_text="later")

    assert len(log.query()) == 2
    assert len(log.query(user_id="u1")) == 1
    assert len(log.query(user_id="u2")) == 1
    assert len(log.query(user_id="u3")) == 0
    teardown()


def test_chain_integrity():
    setup()
    log = AuditLog(AUDIT_FILE)
    log.record(provider="test", input_text="one", output_text="1")
    log.record(provider="test", input_text="two", output_text="2")
    log.record(provider="test", input_text="three", output_text="3")

    valid, count = log.verify_chain()
    assert valid is True
    assert count == 3
    teardown()


def test_pii_auto_detection():
    setup()
    log = AuditLog(AUDIT_FILE)
    entry = log.record(
        provider="test",
        input_text="Email me at test@example.com",
        output_text="ok",
    )
    assert "EMAIL" in entry.pii_detected
    assert "[EMAIL_REDACTED]" in entry.input_redacted
    teardown()


def test_stats():
    setup()
    log = AuditLog(AUDIT_FILE)
    log.record(provider="openai", purpose="support", input_text="hi",
               output_text="hello", tokens_in=5, tokens_out=3)
    log.record(provider="openai", purpose="analysis", input_text="data",
               output_text="result", tokens_in=10, tokens_out=20)

    stats = log.stats()
    assert stats["total_entries"] == 2
    assert stats["providers"]["openai"] == 2
    assert stats["total_tokens_in"] == 15
    assert stats["total_tokens_out"] == 23
    teardown()


def test_export_for_user():
    setup()
    log = AuditLog(AUDIT_FILE)
    log.record(provider="test", user_id="target", input_text="a", output_text="b")
    log.record(provider="test", user_id="other", input_text="c", output_text="d")
    log.record(provider="test", user_id="target", input_text="e", output_text="f")

    export = log.export_for_user("target")
    assert len(export) == 2
    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.")
