viral_script_engine/tests/test_phase10.py

"""Phase 10 tests — A/B Testing Environment Layer."""
import math
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from viral_script_engine.environment.trajectory import Trajectory, TrajectoryType
from viral_script_engine.rewards.contrastive_reward import ContrastiveReward, ContrastiveRewardResult

_SCRIPTS_PATH = str(
    Path(__file__).parent.parent / "data" / "test_scripts" / "scripts.json"
)
_CULTURAL_KB_PATH = str(
    Path(__file__).parent.parent / "data" / "cultural_kb.json"
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_claim(claim_id: str, severity: str, critique_class: str) -> MagicMock:
    c = MagicMock()
    c.claim_id = claim_id
    c.severity = severity
    c.critique_class = critique_class
    c.claim_text = f"Test claim {claim_id}"
    return c


def _make_defender(core_quote: str, flagged: list = None) -> MagicMock:
    d = MagicMock()
    d.core_strength_quote = core_quote
    d.flagged_critic_claims = flagged or []
    return d


def _make_trajectory(
    traj_type: str,
    cumulative: float,
    script: str = "Test script body content here.",
) -> Trajectory:
    return Trajectory(
        trajectory_id=f"test_{traj_type}",
        trajectory_type=traj_type,
        initial_script=script,
        current_script=script,
        cumulative_reward=cumulative,
    )


# ---------------------------------------------------------------------------
# Trajectory: forced first action — CRITIC_FIRST
# ---------------------------------------------------------------------------

class TestTrajectoryForcedActionCriticFirst:
    def setup_method(self):
        self.traj = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.0)

    def test_picks_highest_severity_claim(self):
        claims = [
            _make_claim("C1", "low", "pacing_issue"),
            _make_claim("C2", "high", "hook_weakness"),
            _make_claim("C3", "medium", "cta_buried"),
        ]
        action = self.traj.get_forced_first_action(claims, None)
        # highest severity is C2 (high, hook_weakness → hook_rewrite)
        assert action["action_type"] == "hook_rewrite"
        assert action["critique_claim_id"] == "C2"

    def test_maps_cta_buried_to_cta_placement(self):
        claims = [_make_claim("C1", "high", "cta_buried")]
        action = self.traj.get_forced_first_action(claims, None)
        assert action["action_type"] == "cta_placement"
        assert action["target_section"] == "cta"

    def test_maps_cultural_mismatch_to_cultural_ref_sub(self):
        claims = [_make_claim("C1", "high", "cultural_mismatch")]
        action = self.traj.get_forced_first_action(claims, None)
        assert action["action_type"] == "cultural_ref_sub"

    def test_fallback_when_no_claims(self):
        action = self.traj.get_forced_first_action([], None)
        assert action["action_type"] == "hook_rewrite"
        assert "CRITIC_FIRST" in action["reasoning"] or action["reasoning"]

    def test_reasoning_mentions_critic_first(self):
        claims = [_make_claim("C1", "high", "hook_weakness")]
        action = self.traj.get_forced_first_action(claims, None)
        assert "CRITIC_FIRST" in action["reasoning"]


# ---------------------------------------------------------------------------
# Trajectory: forced first action — DEFENDER_FIRST
# ---------------------------------------------------------------------------

class TestTrajectoryForcedActionDefenderFirst:
    def test_picks_cta_when_core_strength_in_hook(self):
        # Script starts with the core quote → hook is precious
        script = "Why does your phone battery lie? Charge to eighty. Never below twenty."
        traj = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.0, script=script)
        defender = _make_defender(core_quote="Why does your phone battery lie?")
        claims = [_make_claim("C1", "high", "hook_weakness")]

        action = traj.get_forced_first_action(claims, defender)
        assert action["action_type"] == "cta_placement", (
            f"Expected cta_placement when core strength is in hook, got {action['action_type']}"
        )

    def test_picks_hook_rewrite_when_core_strength_in_body(self):
        # Script hook is entirely generic; core quote only appears after the first 100 chars
        filler = "Stop wasting your money on things that do not matter at all. " * 2  # >100 chars
        core = "UNIQUE_CORE_PHRASE_XYZ"
        script = filler + core
        traj = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.0, script=script)
        # Core quote appears after position 100 — NOT in hook
        defender = _make_defender(core_quote=core)
        claims = [_make_claim("C1", "high", "hook_weakness")]

        action = traj.get_forced_first_action(claims, defender)
        assert action["action_type"] == "hook_rewrite", (
            f"Expected hook_rewrite when core is NOT in hook, got {action['action_type']}"
        )

    def test_reasoning_mentions_defender_first(self):
        traj = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.0)
        action = traj.get_forced_first_action([], None)
        assert "DEFENDER_FIRST" in action["reasoning"]

    def test_skips_flagged_claims_in_defender_first(self):
        script = "Body content only. No hook magic here at all whatsoever for testing purposes."
        traj = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.0, script=script)
        defender = _make_defender(
            core_quote="definitely not in the hook portion of this script",
            flagged=["C1"],
        )
        claims = [
            _make_claim("C1", "high", "hook_weakness"),
            _make_claim("C2", "medium", "pacing_issue"),
        ]
        action = traj.get_forced_first_action(claims, defender)
        # C1 is flagged, so should pick C2
        assert action["critique_claim_id"] == "C2"


# ---------------------------------------------------------------------------
# ContrastiveReward
# ---------------------------------------------------------------------------

class TestContrastiveReward:
    def setup_method(self):
        self.cr = ContrastiveReward()

    def test_delta_computed_correctly(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.7)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.5)
        result = self.cr.compute(traj_a, traj_b)
        assert abs(result.delta - 0.2) < 1e-9

    def test_winning_trajectory_is_a_when_a_higher(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.8)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.5)
        result = self.cr.compute(traj_a, traj_b)
        assert result.winning_trajectory == "A"
        assert result.winning_trajectory_type == TrajectoryType.CRITIC_FIRST

    def test_winning_trajectory_is_b_when_b_higher(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.4)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.7)
        result = self.cr.compute(traj_a, traj_b)
        assert result.winning_trajectory == "B"
        assert result.winning_trajectory_type == TrajectoryType.DEFENDER_FIRST

    def test_tie_when_delta_is_zero(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.6)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.6)
        result = self.cr.compute(traj_a, traj_b)
        assert result.winning_trajectory == "tie"

    def test_contrast_bonus_near_zero_when_delta_small(self):
        # delta = 0.01 → tanh(0.01 * 3) * 0.2 ≈ 0.006 — near zero
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.51)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.50)
        result = self.cr.compute(traj_a, traj_b)
        assert abs(result.contrast_bonus) < 0.02, (
            f"contrast_bonus should be near 0 for delta=0.01, got {result.contrast_bonus}"
        )

    def test_contrast_bonus_positive_when_delta_large(self):
        # delta = 0.3 → tanh(0.9) * 0.2 ≈ 0.156 — clearly positive
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.7)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.4)
        result = self.cr.compute(traj_a, traj_b)
        assert result.contrast_bonus > 0.1, (
            f"contrast_bonus should be > 0.1 for delta=0.3, got {result.contrast_bonus}"
        )

    def test_final_reward_clipped_to_0_1_upper(self):
        # Very high cumulative rewards should still clip to 1.0
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 5.0)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.1)
        result = self.cr.compute(traj_a, traj_b)
        assert result.final_reward <= 1.0

    def test_final_reward_clipped_to_0_1_lower(self):
        # Negative cumulative rewards should clip to 0.0
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, -1.0)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, -2.0)
        result = self.cr.compute(traj_a, traj_b)
        assert result.final_reward >= 0.0

    def test_final_reward_always_in_0_1(self):
        for cum_a, cum_b in [(0.3, 0.3), (0.9, 0.1), (0.0, 0.0), (0.5, 0.5), (1.0, 0.0)]:
            traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, cum_a)
            traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, cum_b)
            result = self.cr.compute(traj_a, traj_b)
            assert 0.0 <= result.final_reward <= 1.0, (
                f"final_reward={result.final_reward} out of [0,1] for "
                f"cum_a={cum_a}, cum_b={cum_b}"
            )

    def test_base_reward_is_max(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.7)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.5)
        result = self.cr.compute(traj_a, traj_b)
        assert abs(result.base_reward - 0.7) < 1e-9

    def test_result_is_contrastive_reward_result_instance(self):
        traj_a = _make_trajectory(TrajectoryType.CRITIC_FIRST, 0.6)
        traj_b = _make_trajectory(TrajectoryType.DEFENDER_FIRST, 0.4)
        result = self.cr.compute(traj_a, traj_b)
        assert isinstance(result, ContrastiveRewardResult)


# ---------------------------------------------------------------------------
# ABScriptEnv — integration tests using mocked env.step() and reset()
# ---------------------------------------------------------------------------

def _fake_obs(script: str = "Test script.", reward: float = 0.5) -> dict:
    return {
        "current_script": script,
        "original_script": script,
        "region": "pan_india_english",
        "platform": "Reels",
        "niche": "personal finance",
        "step_num": 1,
        "max_steps": 3,
        "debate_history": [],
        "reward_components": {"r1_hook_strength": reward, "total": reward},
        "difficulty_level": "easy",
        "episode_id": "ep_test",
        "current_moderation_flags": [],
        "current_originality_flags": [],
        "creator_profile": None,
    }


def _fake_step_result(script: str = "Test script.", reward: float = 0.5, done: bool = False):
    obs = _fake_obs(script, reward)
    info = {"reward_components": {"r1_hook_strength": reward, "total": reward}}
    return obs, reward, done, False, info


def _make_real_critique(claim_id="C1", severity="high", critique_class="hook_weakness"):
    """Return a MagicMock with real CritiqueClaim objects so pydantic validation passes."""
    from viral_script_engine.agents.critic import CritiqueClaim, CritiqueOutput
    claim = CritiqueClaim(
        claim_id=claim_id,
        critique_class=critique_class,
        claim_text=f"Test {critique_class} claim",
        timestamp_range="0:00-0:05",
        evidence="test evidence here",
        is_falsifiable=True,
        severity=severity,
    )
    mock_crit = MagicMock()
    mock_crit.claims = [claim]
    mock_crit.overall_severity = severity
    return mock_crit


def _make_real_defender(core_quote="hook content here"):
    from viral_script_engine.agents.defender import DefenderOutput
    return DefenderOutput(
        core_strength="Strong hook",
        core_strength_quote=core_quote,
        defense_argument="Preserve this element.",
        flagged_critic_claims=[],
        regional_voice_elements=[],
    )


class TestABScriptEnvMocked:
    """Test ABScriptEnv behaviour with env.step() mocked at the env level."""

    def _make_ab_env(self):
        from viral_script_engine.environment.ab_env import ABScriptEnv
        return ABScriptEnv(
            scripts_path=_SCRIPTS_PATH,
            cultural_kb_path=_CULTURAL_KB_PATH,
            max_steps=3,
            difficulty="easy",
        )

    def _reset_with_mocks(self, ab_env, core_quote="body content deep here", seed=42):
        """
        Reset ab_env with mocked critic, defender, and step calls.
        Uses real CritiqueClaim/DefenderOutput to pass pydantic validation.
        Returns the state dict.
        """
        mock_critique = _make_real_critique("C1", "high", "hook_weakness")
        mock_defender = _make_real_defender(core_quote)

        with patch.object(ab_env.env_a.critic, "critique", return_value=mock_critique), \
             patch.object(ab_env.env_a.defender, "defend", return_value=mock_defender), \
             patch.object(ab_env.env_b.step, "__call__", side_effect=None) if False else \
             patch.object(ab_env.env_a, "step",
                          side_effect=lambda action, **kw: _fake_step_result("Script A.", 0.65)), \
             patch.object(ab_env.env_b, "step",
                          side_effect=lambda action, **kw: _fake_step_result("Script B.", 0.55)):
            state = ab_env.reset(seed=seed)
        return state

    def test_reset_gives_both_trajectory_states(self):
        ab_env = self._make_ab_env()
        state = self._reset_with_mocks(ab_env)

        assert "trajectory_a" in state
        assert "trajectory_b" in state
        assert "delta" in state
        assert "leading_trajectory" in state
        assert "episode_id" in state

    def test_both_envs_start_from_same_script(self):
        ab_env = self._make_ab_env()
        self._reset_with_mocks(ab_env, seed=42)

        # Both trajectories must share the same initial_script (same reset seed)
        assert ab_env._traj_a.initial_script == ab_env._traj_b.initial_script

    def test_step_1_forced_actions_differ(self):
        """
        Traj A (critic_first, hook_weakness claim) → hook_rewrite.
        Traj B (defender_first, core in hook) → cta_placement.
        """
        import json as _json
        scripts = _json.loads(open(_SCRIPTS_PATH).read())
        easy_script = next(s for s in scripts if s["script_id"] == "S01")
        # Use first 30 chars of the real script as the "core quote" so it appears in the hook
        hook_text = easy_script["script_text"][:30]

        ab_env = self._make_ab_env()
        mock_critique = _make_real_critique("C1", "high", "hook_weakness")
        mock_defender = _make_real_defender(core_quote=hook_text)

        with patch.object(ab_env.env_a.critic, "critique", return_value=mock_critique), \
             patch.object(ab_env.env_a.defender, "defend", return_value=mock_defender), \
             patch.object(ab_env.env_a, "step",
                          side_effect=lambda action, **kw: _fake_step_result()), \
             patch.object(ab_env.env_b, "step",
                          side_effect=lambda action, **kw: _fake_step_result()):
            ab_env.reset(seed=42)

        action_a = ab_env._forced_action_a.get("action_type")
        action_b = ab_env._forced_action_b.get("action_type")
        assert action_a == "hook_rewrite", (
            f"CRITIC_FIRST: expected hook_rewrite, got {action_a}"
        )
        assert action_b == "cta_placement", (
            f"DEFENDER_FIRST (core in hook): expected cta_placement, got {action_b}"
        )

    def test_step_applies_same_action_to_both(self):
        ab_env = self._make_ab_env()
        self._reset_with_mocks(ab_env)

        step_calls_a: list = []
        step_calls_b: list = []

        def track_a(action, **kw):
            step_calls_a.append(action)
            return _fake_step_result("A after free step", 0.7, done=True)

        def track_b(action, **kw):
            step_calls_b.append(action)
            return _fake_step_result("B after free step", 0.6, done=True)

        free_action = {
            "action_type": "cta_placement",
            "target_section": "cta",
            "instruction": "Move CTA to end.",
            "critique_claim_id": "C1",
            "reasoning": "test",
        }
        with patch.object(ab_env.env_a, "step", side_effect=track_a), \
             patch.object(ab_env.env_b, "step", side_effect=track_b):
            ab_env.step(free_action)

        assert len(step_calls_a) == 1
        assert len(step_calls_b) == 1
        assert step_calls_a[0]["action_type"] == step_calls_b[0]["action_type"] == "cta_placement"

    def test_state_returns_correct_delta(self):
        ab_env = self._make_ab_env()
        self._reset_with_mocks(ab_env)

        # Manually set cumulative rewards to known values
        ab_env._traj_a.cumulative_reward = 0.7
        ab_env._traj_b.cumulative_reward = 0.5

        state = ab_env.state()
        assert abs(state["delta"] - 0.2) < 1e-9
        assert state["leading_trajectory"] == "A"
        assert "trajectory_a" in state
        assert "trajectory_b" in state