"""
Empirical test suite for ai_assert.py — Stochastic generator tests.

This fills the AX52 zero-dimension: all 64 existing tests use deterministic
mock generators, providing zero empirical coverage of actual stochastic behavior.

Architecture:
  StochasticGenerator — controlled non-deterministic output simulator
  5 independent test classes (KV7 — no shared state):
    1. TestStochasticGeneratorBehavior — verify the generator itself
    2. TestRetryConvergence — statistical retry/feedback dynamics
    3. TestFeedbackImpact — feedback improves outcomes vs. blind retry
    4. TestStochasticEdgeCases — edge cases under randomness
    5. TestStatisticalProperties — aggregate properties across N runs

Design constraints:
  - Fixed seeds for reproducible CI runs (deterministic-under-seed)
  - Statistical tests use N=50 runs with conservative thresholds
  - All scores bounded < 1.0 (T6 convergence bound)
  - Multiplicative gate (AX52) verified under stochastic conditions
"""

import unittest
import random
import json
import sys
import os

# Ensure the project root is on the path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ai_assert import (
    CheckResult,
    Constraint,
    AiAssertResult,
    ai_assert,
    reliable,
    max_length,
    min_length,
    contains,
    not_contains,
    valid_json,
    matches_schema,
    custom,
)


# ===========================================================================
# STOCHASTIC GENERATOR — Controlled Non-Deterministic Output
# ===========================================================================

class StochasticGenerator:
    """
    Simulates non-deterministic AI output with controllable parameters.

    - target: the "correct" output the generator is trying to produce
    - noise_rate: probability of corruption on first call (0.0 = always correct)
    - improvement_rate: multiplicative decay of noise per attempt (0.0 = instant fix)
    - seed: for reproducible runs; None = truly random
    - feedback_sensitivity: how much feedback reduces noise (0.0 = ignores feedback)

    On each call:
      p_correct = 1 - noise_rate * (improvement_rate ^ (call_count - 1))
      If feedback is detected in the prompt, p_correct gets a boost.
      If p_correct > random(), return target; else return corrupted version.
    """

    def __init__(
        self,
        target: str,
        noise_rate: float = 0.5,
        improvement_rate: float = 0.5,
        seed: int | None = None,
        feedback_sensitivity: float = 0.3,
        corruption_fn=None,
    ):
        self.target = target
        self.noise_rate = max(0.0, min(noise_rate, 1.0))
        self.improvement_rate = max(0.0, min(improvement_rate, 1.0))
        self.feedback_sensitivity = max(0.0, min(feedback_sensitivity, 1.0))
        self.rng = random.Random(seed)
        self.call_count = 0
        self._corruption_fn = corruption_fn or self._default_corrupt

    def __call__(self, prompt: str) -> str:
        self.call_count += 1

        # Base probability of producing correct output
        p_correct = 1.0 - self.noise_rate * (
            self.improvement_rate ** (self.call_count - 1)
        )

        # Feedback boost: if the prompt contains correction markers, boost probability
        if "CORRECTION NEEDED" in prompt or "fix" in prompt.lower():
            p_correct = min(p_correct + self.feedback_sensitivity, 0.9999)

        if self.rng.random() < p_correct:
            return self.target
        return self._corruption_fn(self.target, self.rng)

    def reset(self):
        """Reset call count for reuse across test runs."""
        self.call_count = 0

    @staticmethod
    def _default_corrupt(target: str, rng: random.Random) -> str:
        """Default corruption: truncate, shuffle, or substitute."""
        strategy = rng.choice(["truncate", "shuffle", "substitute", "empty"])
        if strategy == "truncate":
            cut = rng.randint(0, max(len(target) // 2, 1))
            return target[:cut]
        elif strategy == "shuffle":
            chars = list(target)
            rng.shuffle(chars)
            return "".join(chars)
        elif strategy == "substitute":
            return "WRONG_OUTPUT_" + str(rng.randint(0, 999))
        else:
            return ""

    @staticmethod
    def json_corrupt(target: str, rng: random.Random) -> str:
        """Corruption function for JSON targets."""
        strategy = rng.choice(["invalid", "missing_key", "wrong_type", "empty"])
        if strategy == "invalid":
            return "not valid json {{{ broken"
        elif strategy == "missing_key":
            try:
                data = json.loads(target)
                if isinstance(data, dict) and data:
                    key = rng.choice(list(data.keys()))
                    del data[key]
                    return json.dumps(data)
            except (json.JSONDecodeError, IndexError):
                pass
            return "{}"
        elif strategy == "wrong_type":
            return "[]"  # Array instead of object
        else:
            return ""


# ===========================================================================
# TEST CLASS 1: StochasticGenerator Behavior
# ===========================================================================
class TestStochasticGeneratorBehavior(unittest.TestCase):
    """Verify the stochastic generator itself behaves as specified."""

    def test_zero_noise_always_correct(self):
        """noise_rate=0 should always produce correct output."""
        gen = StochasticGenerator("hello world", noise_rate=0.0, seed=42)
        for _ in range(20):
            gen.reset()
            self.assertEqual(gen("test"), "hello world")

    def test_full_noise_first_call_usually_corrupted(self):
        """noise_rate=1.0 should almost always corrupt the first call."""
        corrupt_count = 0
        for seed in range(50):
            gen = StochasticGenerator("hello", noise_rate=1.0, improvement_rate=0.9, seed=seed)
            output = gen("test")
            if output != "hello":
                corrupt_count += 1
        # With noise_rate=1.0 and improvement_rate=0.9, first call has 10% chance correct
        # So ~90% should be corrupted. Allow some slack.
        self.assertGreater(corrupt_count, 30,
                           f"Expected most first calls corrupted, got {corrupt_count}/50")

    def test_improvement_over_calls(self):
        """Each successive call should have higher probability of correct output."""
        gen = StochasticGenerator("target", noise_rate=0.8, improvement_rate=0.5, seed=42)
        # Call many times. Later calls should trend correct.
        results = [gen("test") for _ in range(10)]
        # Last 5 should have more correct than first 5 (statistically)
        early_correct = sum(1 for r in results[:5] if r == "target")
        late_correct = sum(1 for r in results[5:] if r == "target")
        # This specific seed should show improvement, but we check >=
        self.assertGreaterEqual(late_correct, early_correct,
                                "Later calls should trend toward correct output")

    def test_seed_reproducibility(self):
        """Same seed must produce identical sequence."""
        gen1 = StochasticGenerator("hello", noise_rate=0.5, seed=123)
        gen2 = StochasticGenerator("hello", noise_rate=0.5, seed=123)
        seq1 = [gen1("test") for _ in range(10)]
        seq2 = [gen2("test") for _ in range(10)]
        self.assertEqual(seq1, seq2, "Same seed must produce identical outputs")

    def test_different_seeds_differ(self):
        """Different seeds should (very likely) produce different sequences."""
        gen1 = StochasticGenerator("hello", noise_rate=0.5, seed=1)
        gen2 = StochasticGenerator("hello", noise_rate=0.5, seed=999)
        seq1 = [gen1("test") for _ in range(10)]
        seq2 = [gen2("test") for _ in range(10)]
        self.assertNotEqual(seq1, seq2,
                            "Different seeds should produce different outputs")

    def test_reset_restarts_call_count(self):
        """reset() should bring call_count back to 0."""
        gen = StochasticGenerator("hello", noise_rate=0.5, seed=42)
        gen("test")
        gen("test")
        self.assertEqual(gen.call_count, 2)
        gen.reset()
        self.assertEqual(gen.call_count, 0)

    def test_feedback_sensitivity_boost(self):
        """Prompts containing 'CORRECTION NEEDED' should boost correct probability."""
        correct_without_feedback = 0
        correct_with_feedback = 0
        for seed in range(100):
            gen1 = StochasticGenerator("target", noise_rate=0.7, seed=seed,
                                       feedback_sensitivity=0.5)
            gen2 = StochasticGenerator("target", noise_rate=0.7, seed=seed,
                                       feedback_sensitivity=0.5)
            out1 = gen1("plain prompt")
            out2 = gen2("CORRECTION NEEDED: fix the output")
            if out1 == "target":
                correct_without_feedback += 1
            if out2 == "target":
                correct_with_feedback += 1
        self.assertGreater(correct_with_feedback, correct_without_feedback,
                           f"Feedback should increase correct rate: "
                           f"{correct_with_feedback} vs {correct_without_feedback}")

    def test_custom_corruption_fn(self):
        """Custom corruption function should be used instead of default."""
        def always_corrupt(target, rng):
            return "CUSTOM_CORRUPT"

        gen = StochasticGenerator("hello", noise_rate=1.0, improvement_rate=1.0,
                                  seed=42, corruption_fn=always_corrupt)
        output = gen("test")
        self.assertEqual(output, "CUSTOM_CORRUPT")


# ===========================================================================
# TEST CLASS 2: Retry Convergence Under Stochastic Conditions
# ===========================================================================
class TestRetryConvergence(unittest.TestCase):
    """Verify ai_assert's retry mechanism converges under stochastic generators."""

    def test_eventual_pass_with_moderate_noise(self):
        """With moderate noise, retries should usually succeed."""
        pass_count = 0
        N = 50
        for seed in range(N):
            gen = StochasticGenerator("hello world", noise_rate=0.6,
                                      improvement_rate=0.4, seed=seed)
            result = ai_assert(
                prompt="say hello world",
                constraints=[contains("hello")],
                generate_fn=gen,
                max_retries=5,
            )
            if result.passed:
                pass_count += 1
        # With generous retries and moderate noise, expect high pass rate
        self.assertGreater(pass_count, N * 0.6,
                           f"Expected >60% pass rate, got {pass_count}/{N}")

    def test_retries_reduce_attempt_count_with_improvement(self):
        """With improvement_rate < 1, later runs should need fewer attempts."""
        early_attempts = []
        late_attempts = []
        for seed in range(50):
            gen = StochasticGenerator("target", noise_rate=0.7,
                                      improvement_rate=0.3, seed=seed)
            result = ai_assert(
                prompt="produce target",
                constraints=[contains("target")],
                generate_fn=gen,
                max_retries=5,
            )
            if result.passed:
                if seed < 25:
                    early_attempts.append(result.attempts)
                else:
                    late_attempts.append(result.attempts)
        # Both groups should have results; average attempts should be low
        if early_attempts and late_attempts:
            avg_early = sum(early_attempts) / len(early_attempts)
            avg_late = sum(late_attempts) / len(late_attempts)
            # Just verify both groups complete — improvement_rate helps internally per run
            self.assertLessEqual(avg_early, 6, f"Early avg too high: {avg_early}")
            self.assertLessEqual(avg_late, 6, f"Late avg too high: {avg_late}")

    def test_high_noise_low_retries_mostly_fail(self):
        """noise_rate=0.95 with max_retries=0 should mostly fail."""
        fail_count = 0
        N = 50
        for seed in range(N):
            gen = StochasticGenerator("target", noise_rate=0.95,
                                      improvement_rate=0.99, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("target")],
                generate_fn=gen,
                max_retries=0,
            )
            if not result.passed:
                fail_count += 1
        self.assertGreater(fail_count, N * 0.7,
                           f"Expected >70% fail rate, got {fail_count}/{N}")

    def test_zero_noise_always_passes(self):
        """noise_rate=0.0 should pass on first attempt every time."""
        for seed in range(20):
            gen = StochasticGenerator("hello world", noise_rate=0.0, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("hello")],
                generate_fn=gen,
                max_retries=3,
            )
            self.assertTrue(result.passed)
            self.assertEqual(result.attempts, 1)

    def test_multiple_constraints_stochastic(self):
        """Multiple constraints under stochastic conditions."""
        pass_count = 0
        N = 50
        target = "hello world"
        for seed in range(N):
            gen = StochasticGenerator(target, noise_rate=0.5,
                                      improvement_rate=0.3, seed=seed)
            result = ai_assert(
                prompt="say hello world",
                constraints=[
                    contains("hello"),
                    min_length(5),
                    max_length(100),
                ],
                generate_fn=gen,
                max_retries=5,
            )
            if result.passed:
                pass_count += 1
        self.assertGreater(pass_count, N * 0.5,
                           f"Expected >50% pass with multiple constraints, got {pass_count}/{N}")

    def test_history_length_matches_attempts(self):
        """History should always have exactly `attempts` entries."""
        for seed in range(20):
            gen = StochasticGenerator("target", noise_rate=0.5, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("target")],
                generate_fn=gen,
                max_retries=3,
            )
            self.assertEqual(len(result.history), result.attempts,
                             f"Seed {seed}: history length mismatch")


# ===========================================================================
# TEST CLASS 3: Feedback Impact
# ===========================================================================
class TestFeedbackImpact(unittest.TestCase):
    """Verify that feedback (default or custom) improves outcomes vs. blind retry."""

    def test_default_feedback_improves_pass_rate(self):
        """Default feedback mechanism should improve pass rate over no-feedback."""
        pass_with_feedback = 0
        pass_without_feedback = 0
        N = 50

        for seed in range(N):
            # WITH feedback (default): ai_assert builds correction prompts
            gen1 = StochasticGenerator("hello world", noise_rate=0.7,
                                       improvement_rate=0.6, seed=seed,
                                       feedback_sensitivity=0.4)
            result1 = ai_assert(
                prompt="say hello world",
                constraints=[contains("hello")],
                generate_fn=gen1,
                max_retries=3,
            )
            if result1.passed:
                pass_with_feedback += 1

            # WITHOUT feedback: generator ignores feedback markers
            gen2 = StochasticGenerator("hello world", noise_rate=0.7,
                                       improvement_rate=0.6, seed=seed,
                                       feedback_sensitivity=0.0)  # ignores feedback
            result2 = ai_assert(
                prompt="say hello world",
                constraints=[contains("hello")],
                generate_fn=gen2,
                max_retries=3,
            )
            if result2.passed:
                pass_without_feedback += 1

        # Feedback-sensitive generator should have higher pass rate
        self.assertGreaterEqual(pass_with_feedback, pass_without_feedback,
                                f"Feedback should help: {pass_with_feedback} vs {pass_without_feedback}")

    def test_custom_feedback_fn_called_on_failure(self):
        """Custom feedback function should be invoked when constraints fail."""
        feedback_calls = []

        def track_feedback(output, failed_checks):
            feedback_calls.append((output, len(failed_checks)))
            failures = ", ".join(c.name for c in failed_checks)
            return f"Fix these: {failures}. Previous: {output[:50]}"

        gen = StochasticGenerator("hello world", noise_rate=0.8,
                                  improvement_rate=0.3, seed=42)
        result = ai_assert(
            prompt="say hello",
            constraints=[contains("hello")],
            generate_fn=gen,
            max_retries=5,
            feedback_fn=track_feedback,
        )

        if result.attempts > 1:
            self.assertGreater(len(feedback_calls), 0,
                               "Custom feedback should have been called on retry")

    def test_feedback_contains_failure_info(self):
        """Default feedback prompt should contain specific failure information."""
        prompts_received = []

        def capturing_gen(prompt):
            prompts_received.append(prompt)
            if len(prompts_received) >= 3:
                return "hello world"
            return "wrong"

        result = ai_assert(
            prompt="say hello",
            constraints=[contains("hello")],
            generate_fn=capturing_gen,
            max_retries=5,
        )

        if len(prompts_received) > 1:
            # Second prompt should contain correction info
            self.assertIn("CORRECTION NEEDED", prompts_received[1])
            self.assertIn("hello", prompts_received[1].lower())


# ===========================================================================
# TEST CLASS 4: Stochastic Edge Cases
# ===========================================================================
class TestStochasticEdgeCases(unittest.TestCase):
    """Edge cases unique to stochastic (non-deterministic) execution."""

    def test_json_constraint_with_stochastic_json(self):
        """valid_json constraint under stochastic JSON generation."""
        target = '{"name": "Alice", "age": 30}'
        pass_count = 0
        N = 30
        for seed in range(N):
            gen = StochasticGenerator(
                target, noise_rate=0.5, improvement_rate=0.3,
                seed=seed, corruption_fn=StochasticGenerator.json_corrupt,
            )
            result = ai_assert(
                prompt="generate JSON",
                constraints=[valid_json()],
                generate_fn=gen,
                max_retries=5,
            )
            if result.passed:
                pass_count += 1
        self.assertGreater(pass_count, N * 0.4,
                           f"JSON constraint should pass often enough: {pass_count}/{N}")

    def test_schema_constraint_stochastic(self):
        """matches_schema under stochastic conditions."""
        target = '{"name": "Alice", "age": 30, "email": "a@b.com"}'
        schema = {"required": {"name", "age", "email"}}
        pass_count = 0
        N = 30
        for seed in range(N):
            gen = StochasticGenerator(
                target, noise_rate=0.5, improvement_rate=0.3,
                seed=seed, corruption_fn=StochasticGenerator.json_corrupt,
            )
            result = ai_assert(
                prompt="generate user JSON",
                constraints=[matches_schema(schema)],
                generate_fn=gen,
                max_retries=5,
            )
            if result.passed:
                pass_count += 1
        self.assertGreater(pass_count, N * 0.3,
                           f"Schema constraint should pass sometimes: {pass_count}/{N}")

    def test_contradictory_constraints_never_pass(self):
        """Contradictory constraints should never pass even with retries."""
        for seed in range(20):
            gen = StochasticGenerator("hello", noise_rate=0.0, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("hello"), not_contains("hello")],
                generate_fn=gen,
                max_retries=5,
            )
            self.assertFalse(result.passed,
                             "Contradictory constraints must never pass")

    def test_empty_target_with_min_length(self):
        """Generator targeting empty string should fail min_length consistently."""
        for seed in range(20):
            gen = StochasticGenerator("", noise_rate=0.0, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[min_length(5)],
                generate_fn=gen,
                max_retries=3,
            )
            self.assertFalse(result.passed)

    def test_very_long_target(self):
        """Very long target with max_length constraint under noise."""
        target = "x" * 100
        pass_count = 0
        for seed in range(20):
            gen = StochasticGenerator(target, noise_rate=0.3, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[max_length(200), min_length(50)],
                generate_fn=gen,
                max_retries=3,
            )
            if result.passed:
                pass_count += 1
        self.assertGreater(pass_count, 5,
                           f"Long target should pass sometimes: {pass_count}/20")

    def test_decorator_with_stochastic_generator(self):
        """@reliable decorator should work with stochastic generators."""
        gen = StochasticGenerator("hello world", noise_rate=0.3,
                                  improvement_rate=0.3, seed=42)

        @reliable(constraints=[contains("hello"), max_length(100)], max_retries=5)
        def stochastic_fn(prompt):
            return gen(prompt)

        result = stochastic_fn("say hello")
        self.assertIsInstance(result, AiAssertResult)
        # With low noise and many retries, should usually pass
        # But even if it doesn't, the type and structure must be correct
        self.assertIsInstance(result.passed, bool)
        self.assertGreater(result.attempts, 0)


# ===========================================================================
# TEST CLASS 5: Statistical Properties (Aggregate Behavior)
# ===========================================================================
class TestStatisticalProperties(unittest.TestCase):
    """
    Verify aggregate statistical properties across many stochastic runs.
    These tests establish probabilistic guarantees about ai_assert's behavior.
    """

    def test_composite_score_always_bounded(self):
        """Composite score must be < 1.0 across all stochastic runs (T6)."""
        for seed in range(50):
            gen = StochasticGenerator("hello world", noise_rate=0.3, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("hello"), max_length(100)],
                generate_fn=gen,
                max_retries=3,
            )
            self.assertLess(result.composite_score, 1.0,
                            f"T6 violated at seed {seed}: {result.composite_score}")

    def test_check_scores_always_bounded(self):
        """Individual check scores must be < 1.0 under all conditions (T6)."""
        for seed in range(50):
            gen = StochasticGenerator("hello world", noise_rate=0.3, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[contains("hello"), min_length(3), max_length(100)],
                generate_fn=gen,
                max_retries=3,
            )
            for check in result.checks:
                self.assertLess(check.score, 1.0,
                                f"T6 violated: {check.name} score={check.score} seed={seed}")

    def test_pass_rate_monotonic_with_retries(self):
        """More retries should yield equal or higher pass rate (monotonicity)."""
        N = 50
        results_by_retries = {}
        for max_r in [0, 1, 3, 5]:
            pass_count = 0
            for seed in range(N):
                gen = StochasticGenerator("target", noise_rate=0.6,
                                          improvement_rate=0.4, seed=seed)
                result = ai_assert(
                    prompt="test",
                    constraints=[contains("target")],
                    generate_fn=gen,
                    max_retries=max_r,
                )
                if result.passed:
                    pass_count += 1
            results_by_retries[max_r] = pass_count

        # Verify monotonicity: more retries → equal or more passes
        retry_values = sorted(results_by_retries.keys())
        for i in range(len(retry_values) - 1):
            r1, r2 = retry_values[i], retry_values[i + 1]
            self.assertGreaterEqual(
                results_by_retries[r2], results_by_retries[r1],
                f"Monotonicity violated: retries={r1}→{results_by_retries[r1]}, "
                f"retries={r2}→{results_by_retries[r2]}"
            )

    def test_multiplicative_gate_under_noise(self):
        """AX52: If one constraint is unsatisfiable, all runs should fail."""
        target = "hello world"
        for seed in range(20):
            gen = StochasticGenerator(target, noise_rate=0.0, seed=seed)
            result = ai_assert(
                prompt="test",
                constraints=[
                    contains("hello"),          # always passes (noise=0)
                    contains("IMPOSSIBLE_XYZ"),  # always fails
                ],
                generate_fn=gen,
                max_retries=5,
            )
            self.assertFalse(result.passed,
                             f"AX52: Unsatisfiable constraint must collapse gate (seed {seed})")

    def test_average_attempts_decrease_with_lower_noise(self):
        """Lower noise should mean fewer average attempts needed."""
        N = 50
        avg_attempts = {}
        for noise in [0.2, 0.5, 0.8]:
            total_attempts = 0
            for seed in range(N):
                gen = StochasticGenerator("target", noise_rate=noise,
                                          improvement_rate=0.3, seed=seed)
                result = ai_assert(
                    prompt="test",
                    constraints=[contains("target")],
                    generate_fn=gen,
                    max_retries=5,
                )
                total_attempts += result.attempts
            avg_attempts[noise] = total_attempts / N

        # Lower noise should need fewer attempts on average
        self.assertLessEqual(avg_attempts[0.2], avg_attempts[0.5],
                             f"Lower noise should need fewer attempts: "
                             f"0.2→{avg_attempts[0.2]:.1f}, 0.5→{avg_attempts[0.5]:.1f}")
        self.assertLessEqual(avg_attempts[0.5], avg_attempts[0.8],
                             f"Lower noise should need fewer attempts: "
                             f"0.5→{avg_attempts[0.5]:.1f}, 0.8→{avg_attempts[0.8]:.1f}")

    def test_no_side_effects_between_runs(self):
        """KV7: Independent runs must not affect each other."""
        gen1 = StochasticGenerator("hello", noise_rate=0.3, seed=42)
        gen2 = StochasticGenerator("hello", noise_rate=0.3, seed=42)

        # Run gen1 through ai_assert
        result1 = ai_assert(
            prompt="test",
            constraints=[contains("hello")],
            generate_fn=gen1,
            max_retries=3,
        )

        # Run gen2 through ai_assert (same seed, fresh instance)
        result2 = ai_assert(
            prompt="test",
            constraints=[contains("hello")],
            generate_fn=gen2,
            max_retries=3,
        )

        self.assertEqual(result1.passed, result2.passed, "KV7: Same seed should give same result")
        self.assertEqual(result1.attempts, result2.attempts, "KV7: Same seed should give same attempts")
        self.assertEqual(result1.output, result2.output, "KV7: Same seed should give same output")


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