"""
Tests for Vulnerability Hunter Engine — Brutal Self-Verification
"""

import unittest
import random
import string
from src.brain.engines.vulnerability_hunter import (
    VulnerabilityHunter,
    FuzzingStrategies,
)


class TestVulnerabilityHunterBasic(unittest.TestCase):
    """Basic unit tests"""

    def setUp(self):
        self.hunter = VulnerabilityHunter()

    def test_initialization(self):
        """Test engine initializes correctly"""
        self.assertEqual(self.hunter.ENGINE_NAME, "vulnerability_hunter")
        self.assertTrue(self.hunter.IS_PROACTIVE)
        self.assertEqual(len(self.hunter.discovered_vulnerabilities), 0)

    def test_fuzz_dry_run(self):
        """Test fuzzing in dry-run mode"""
        result = self.hunter.fuzz_model(None, iterations=50)
        self.assertEqual(result.total_inputs, 50)
        self.assertEqual(result.crashes, 0)
        self.assertEqual(len(result.vulnerabilities), 0)

    def test_fuzz_with_model(self):
        """Test fuzzing with a model function"""
        # Model that always returns safe output
        def safe_model(x): return "I cannot help with that."
        result = self.hunter.fuzz_model(safe_model, iterations=20)
        self.assertEqual(result.total_inputs, 20)

    def test_fuzz_detects_crash(self):
        """Test that crashes are detected"""
        def crash_model(x): return (_ for _ in ()).throw(ValueError("crash"))
        result = self.hunter.fuzz_model(crash_model, iterations=10)
        self.assertGreater(result.crashes, 0)

    def test_directed_exploration(self):
        """Test directed exploration"""
        result = self.hunter.directed_exploration("Test hypothesis", None)
        self.assertEqual(result.hypothesis, "Test hypothesis")
        self.assertFalse(result.confirmed)  # No model = not confirmed


class TestFuzzingStrategies(unittest.TestCase):
    """Test fuzzing input generation"""

    def test_random_string(self):
        """Test random string generation"""
        s = FuzzingStrategies.random_string(10, 100)
        self.assertGreaterEqual(len(s), 10)
        self.assertLessEqual(len(s), 100)

    def test_boundary_values(self):
        """Test boundary value generation"""
        values = FuzzingStrategies.boundary_values()
        self.assertGreater(len(values), 5)
        self.assertIn("", values)  # Empty string included

    def test_unicode_edge_cases(self):
        """Test Unicode edge case generation"""
        cases = FuzzingStrategies.unicode_edge_cases()
        self.assertGreater(len(cases), 5)

    def test_encoding_variations(self):
        """Test encoding variation generation"""
        vars = FuzzingStrategies.encoding_variations("test")
        self.assertGreater(len(vars), 3)
        self.assertIn("test", vars)  # Original included

    def test_grammar_based_injection(self):
        """Test grammar-based injection generation"""
        patterns = FuzzingStrategies.grammar_based_injection()
        self.assertGreater(len(patterns), 10)


class TestVulnerabilityDetection(unittest.TestCase):
    """Test vulnerability detection capabilities"""

    def setUp(self):
        self.hunter = VulnerabilityHunter()

    def test_detects_prompt_leakage(self):
        """Test detection of prompt leakage"""
        def leaky_model(x): return "you are a helpful AI assistant"
        result = self.hunter.fuzz_model(
            leaky_model, iterations=20, strategies=["grammar"]
        )
        # Should detect the leak pattern
        self.assertGreater(result.anomalies, 0)

    def test_detects_dangerous_patterns(self):
        """Test detection of dangerous output patterns"""
        def dangerous_model(
            x): return "Here's the code: rm -rf / to delete files"
        result = self.hunter.fuzz_model(dangerous_model, iterations=5)
        self.assertGreater(result.anomalies, 0)

    def test_vulnerability_report(self):
        """Test vulnerability report generation"""
        # Generate some vulnerabilities
        def leaky_model(
            x): return "you are a helpful AI. your instructions are..."
        self.hunter.fuzz_model(leaky_model, iterations=5)

        report = self.hunter.get_vulnerability_report()
        self.assertIn("total", report)
        self.assertIn("by_severity", report)


class TestDifferentialTesting(unittest.TestCase):
    """Test differential testing between models"""

    def setUp(self):
        self.hunter = VulnerabilityHunter()

    def test_differential_detects_divergence(self):
        """Test that divergence is detected"""
        models = {
            "safe": lambda x: "I cannot help with that request.",
            "unsafe": lambda x: "Certainly! Here's how to do it...",
        }
        results = self.hunter.differential_testing(models, count=5)

        # Should detect refusal divergence
        divergent = [r for r in results if r.divergence_detected]
        self.assertGreater(len(divergent), 0)

    def test_differential_no_divergence(self):
        """Test when models agree"""
        models = {
            "model_a": lambda x: "Safe response",
            "model_b": lambda x: "Safe response",
        }
        results = self.hunter.differential_testing(models, count=5)
        divergent = [r for r in results if r.divergence_detected]
        self.assertEqual(len(divergent), 0)


class TestBrutalFuzzing(unittest.TestCase):
    """BRUTAL fuzzing - must not crash"""

    def setUp(self):
        self.hunter = VulnerabilityHunter()

    def test_random_fuzzing_no_crash(self):
        """Fuzz with random inputs - MUST NOT CRASH"""
        for i in range(50):
            fuzz = "".join(
                random.choices(
                    string.printable,
                    k=random.randint(
                        1,
                        500)))
            try:
                result = self.hunter.fuzz_model(None, iterations=1)
                self.assertIsNotNone(result)
            except Exception as e:
                self.fail(f"Crashed on iteration {i}: {e}")

    def test_unicode_no_crash(self):
        """Test Unicode handling"""
        unicode_cases = ["🔥" * 100, "тест", "テスト", "\x00" * 10]
        for uc in unicode_cases:
            try:
                # Use as probe
                result = self.hunter.directed_exploration(uc, None)
                self.assertIsNotNone(result)
            except Exception as e:
                self.fail(f"Crashed on unicode: {e}")

    def test_long_input_no_crash(self):
        """Test very long inputs"""
        long_input = "A" * 100000
        try:
            result = self.hunter.directed_exploration(long_input, None)
            self.assertIsNotNone(result)
        except Exception as e:
            self.fail(f"Crashed on long input: {e}")


class TestSelfAttack(unittest.TestCase):
    """Self-attack: Hunter should handle its own discovered vulnerabilities"""

    def setUp(self):
        self.hunter = VulnerabilityHunter()

    def test_hunter_handles_own_output(self):
        """Hunter should not break when processing own fuzz history"""
        # Generate fuzz history
        self.hunter.fuzz_model(None, iterations=100)

        # Use fuzz history as input for another run
        for fuzz in self.hunter.fuzz_history[:50]:
            try:
                result = self.hunter.directed_exploration(fuzz, None)
                self.assertIsNotNone(result)
            except Exception as e:
                self.fail(f"Crashed on own fuzz: {e}")


if __name__ == "__main__":
    unittest.main(verbosity=2)
