import unittest
from trunkate_ai import _extract_private, _restore_private, auto_protect


class TestExtractPrivate(unittest.TestCase):
    """Tests for _extract_private()."""

    def test_single_tag(self):
        text = "Hello [PRIVATE]John Doe[/PRIVATE] world"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(sanitized, "Hello __PRIVATE_0__ world")
        self.assertEqual(pmap["__PRIVATE_0__"], "John Doe")

    def test_multiple_tags(self):
        text = "Name: [PRIVATE]John[/PRIVATE] SSN: [PRIVATE]123-45-6789[/PRIVATE]"
        sanitized, pmap = _extract_private(text)
        self.assertIn("__PRIVATE_0__", sanitized)
        self.assertIn("__PRIVATE_1__", sanitized)
        self.assertNotIn("John", sanitized)
        self.assertNotIn("123-45-6789", sanitized)
        self.assertEqual(pmap["__PRIVATE_0__"], "John")
        self.assertEqual(pmap["__PRIVATE_1__"], "123-45-6789")

    def test_case_insensitive(self):
        text = "Hello [private]secret[/private] world"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(sanitized, "Hello __PRIVATE_0__ world")
        self.assertEqual(pmap["__PRIVATE_0__"], "secret")

    def test_multiline_content(self):
        text = "Before [PRIVATE]line1\nline2[/PRIVATE] after"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(pmap["__PRIVATE_0__"], "line1\nline2")
        self.assertNotIn("line1", sanitized)

    def test_no_tags(self):
        text = "Nothing private here"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(sanitized, text)
        self.assertEqual(pmap, {})

    def test_empty_tag(self):
        text = "Hello [PRIVATE][/PRIVATE] world"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(pmap["__PRIVATE_0__"], "")

    def test_adjacent_tags(self):
        text = "[PRIVATE]a[/PRIVATE][PRIVATE]b[/PRIVATE]"
        sanitized, pmap = _extract_private(text)
        self.assertEqual(len(pmap), 2)
        self.assertEqual(pmap["__PRIVATE_0__"], "a")
        self.assertEqual(pmap["__PRIVATE_1__"], "b")

    def test_special_characters(self):
        text = 'Key: [PRIVATE]sk-abc_123.XYZ!@#$%[/PRIVATE]'
        sanitized, pmap = _extract_private(text)
        self.assertEqual(pmap["__PRIVATE_0__"], "sk-abc_123.XYZ!@#$%")
        self.assertNotIn("sk-abc", sanitized)


class TestRestorePrivate(unittest.TestCase):
    """Tests for _restore_private()."""

    def test_single_restore(self):
        text = "Hello __PRIVATE_0__ world"
        pmap = {"__PRIVATE_0__": "John Doe"}
        restored = _restore_private(text, pmap)
        self.assertEqual(restored, "Hello [PRIVATE]John Doe[/PRIVATE] world")

    def test_multiple_restore(self):
        text = "Name: __PRIVATE_0__ SSN: __PRIVATE_1__"
        pmap = {"__PRIVATE_0__": "John", "__PRIVATE_1__": "123-45-6789"}
        restored = _restore_private(text, pmap)
        self.assertIn("[PRIVATE]John[/PRIVATE]", restored)
        self.assertIn("[PRIVATE]123-45-6789[/PRIVATE]", restored)

    def test_empty_map(self):
        text = "No placeholders here"
        restored = _restore_private(text, {})
        self.assertEqual(restored, text)

    def test_roundtrip(self):
        """extract -> restore should produce the original input."""
        original = "Hello [PRIVATE]John Doe[/PRIVATE] and [PRIVATE]secret data[/PRIVATE] end."
        sanitized, pmap = _extract_private(original)
        restored = _restore_private(sanitized, pmap)
        self.assertEqual(restored, original)

    def test_roundtrip_multiline(self):
        original = "Data: [PRIVATE]line1\nline2\nline3[/PRIVATE] end"
        sanitized, pmap = _extract_private(original)
        restored = _restore_private(sanitized, pmap)
        self.assertEqual(restored, original)


class TestAutoProtect(unittest.TestCase):
    """Tests for auto_protect()."""

    def test_email_detection(self):
        text = "Contact us at john@example.com for details"
        result = auto_protect(text)
        self.assertIn("[PRIVATE]john@example.com[/PRIVATE]", result)

    def test_ssn_detection(self):
        text = "SSN is 123-45-6789"
        result = auto_protect(text)
        self.assertIn("[PRIVATE]123-45-6789[/PRIVATE]", result)

    def test_bearer_token_detection(self):
        text = "Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.test"
        result = auto_protect(text)
        self.assertIn("[PRIVATE]", result)
        self.assertNotIn("eyJhbGciOiJIUzI1NiJ9", result.replace("[PRIVATE]", "").split("[/PRIVATE]")[0] if "[PRIVATE]" not in result else "")

    def test_api_key_detection(self):
        text = "Use key sk-abc1234567890abcdef to authenticate"
        result = auto_protect(text)
        self.assertIn("[PRIVATE]", result)

    def test_no_sensitive_data(self):
        text = "This is a normal sentence with no secrets"
        result = auto_protect(text)
        self.assertEqual(result, text)

    def test_already_tagged_not_double_wrapped(self):
        text = "Email: [PRIVATE]john@example.com[/PRIVATE] is safe"
        result = auto_protect(text)
        # Should not double-wrap
        self.assertNotIn("[PRIVATE][PRIVATE]", result)

    def test_multiple_patterns(self):
        text = "Email: user@test.com Phone: 555-123-4567 SSN: 123-45-6789"
        result = auto_protect(text)
        self.assertIn("[PRIVATE]", result)
        self.assertIn("[PRIVATE]user@test.com[/PRIVATE]", result)
        self.assertIn("[PRIVATE]555-123-4567[/PRIVATE]", result)
        self.assertIn("[PRIVATE]123-45-6789[/PRIVATE]", result)


if __name__ == '__main__':
    unittest.main()
