viral_script_engine/tests/test_phase2.py

import json
import pytest
from unittest.mock import MagicMock, patch

from viral_script_engine.agents.defender import DefenderAgent, DefenderOutput, DefenderParseError
from viral_script_engine.agents.critic import CritiqueClaim
from viral_script_engine.environment.actions import ActionType, ArbitratorAction
from viral_script_engine.rewards.r3_cultural_alignment import CulturalAlignmentReward
from viral_script_engine.rewards.r4_debate_resolution import DebateResolutionReward, DebateResolutionResult
from viral_script_engine.rewards.r5_defender_preservation import DefenderPreservationReward
from viral_script_engine.rewards.reward_aggregator import RewardAggregator, AntiGamingLog
from viral_script_engine.environment.observations import RewardComponents


# ─── fixtures ────────────────────────────────────────────────────────────────

MOCK_DEFENDER_RESPONSE = json.dumps({
    "core_strength": "Relatable hook about saving money",
    "core_strength_quote": "Let me tell you a secret",
    "defense_argument": "This creates immediate viewer curiosity and should not be changed.",
    "flagged_critic_claims": ["C2"],
    "regional_voice_elements": ["yaar", "ek dum solid"],
})

MOCK_CRITIQUE_CLAIMS = [
    CritiqueClaim(
        claim_id="C1",
        critique_class="hook_weakness",
        claim_text="Weak hook.",
        timestamp_range="0:00-0:03",
        evidence="Let me tell you a secret",
        is_falsifiable=True,
        severity="high",
    ),
    CritiqueClaim(
        claim_id="C2",
        critique_class="cta_buried",
        claim_text="CTA at end.",
        timestamp_range="0:45-0:50",
        evidence="Like and save this video",
        is_falsifiable=True,
        severity="medium",
    ),
]


@pytest.fixture
def mock_defender_llm(monkeypatch):
    monkeypatch.setattr(
        "viral_script_engine.agents.llm_backend.LLMBackend.generate",
        lambda self, sys_prompt, usr_prompt, **kw: MOCK_DEFENDER_RESPONSE,
    )


# ─── Step 1: DefenderAgent ────────────────────────────────────────────────────

def test_defender_parses_output(mock_defender_llm):
    agent = DefenderAgent()
    result = agent.defend(
        script="Let me tell you a secret about saving money. yaar, ek dum solid plan.",
        critic_claims=MOCK_CRITIQUE_CLAIMS,
        region="mumbai_gen_z",
        platform="instagram",
    )
    assert isinstance(result, DefenderOutput)
    assert result.core_strength_quote == "Let me tell you a secret"
    assert "C2" in result.flagged_critic_claims
    assert "yaar" in result.regional_voice_elements


def test_defender_retries_on_invalid_json(monkeypatch):
    call_count = {"n": 0}

    def fake_generate(self, sys_prompt, usr_prompt, **kw):
        call_count["n"] += 1
        if call_count["n"] == 1:
            return "NOT JSON AT ALL"
        return MOCK_DEFENDER_RESPONSE

    monkeypatch.setattr(
        "viral_script_engine.agents.llm_backend.LLMBackend.generate",
        fake_generate,
    )
    agent = DefenderAgent()
    result = agent.defend("script", MOCK_CRITIQUE_CLAIMS, "mumbai_gen_z", "instagram")
    assert isinstance(result, DefenderOutput)
    assert call_count["n"] == 2


def test_defender_raises_after_two_failures(monkeypatch):
    monkeypatch.setattr(
        "viral_script_engine.agents.llm_backend.LLMBackend.generate",
        lambda self, sys_prompt, usr_prompt, **kw: "BAD JSON",
    )
    agent = DefenderAgent()
    with pytest.raises(DefenderParseError):
        agent.defend("script", MOCK_CRITIQUE_CLAIMS, "mumbai_gen_z", "instagram")


# ─── Step 2: R3 CulturalAlignmentReward ──────────────────────────────────────

@pytest.fixture
def r3(tmp_path):
    kb = {
        "mumbai_gen_z": {
            "valid_refs": ["Bandra", "CSMT", "local train", "Swiggy", "IPL"],
            "correct_idioms": ["ek dum solid", "kya scene hai", "full on"],
            "invalid_signals": ["trunk call", "VHS", "walkman"],
            "anachronistic_signals": [],
        },
        "tier2_hindi_belt": {
            "valid_refs": ["kirana store", "sabzi mandi", "jugaad", "panchayat", "mela"],
            "correct_idioms": ["bilkul sahi", "arey bhai", "seedha baat"],
            "invalid_signals": ["SaaS", "venture capital", "coworking space"],
            "anachronistic_signals": [],
        },
    }
    kb_path = tmp_path / "test_kb.json"
    kb_path.write_text(json.dumps(kb), encoding="utf-8")
    return CulturalAlignmentReward(knowledge_base_path=str(kb_path))


def test_r3_scores_regional_script(r3):
    script = "Take the local train to Bandra. IPL is on at night. ek dum solid plan yaar."
    result = r3.score(script, "mumbai_gen_z")
    assert result.score > 0.0
    assert "local train" in result.valid_refs_found or "Bandra" in result.valid_refs_found


def test_r3_scores_non_regional_script_lower(r3):
    script = "Buy on Amazon. Use your credit card. Free delivery available nationwide."
    regional = r3.score(
        "Take local train to Bandra. IPL is on. ek dum solid.", "mumbai_gen_z"
    )
    non_regional = r3.score(script, "mumbai_gen_z")
    assert regional.score >= non_regional.score


def test_r3_penalises_invalid_signals(r3):
    script = "This is like an old VHS walkman trunk call era."
    result = r3.score(script, "mumbai_gen_z")
    assert result.score == 0.0
    assert len(result.invalid_signals_found) > 0


def test_r3_neutral_for_unknown_region(r3):
    result = r3.score("any script", "unknown_region_xyz")
    assert result.score == 0.5


def test_r3_tier2_valid(r3):
    script = "Went to kirana store, met at sabzi mandi, pure jugaad. bilkul sahi plan."
    result = r3.score(script, "tier2_hindi_belt")
    assert result.score > 0.0


def test_r3_tier2_penalises_metro_jargon(r3):
    script = "We raised SaaS venture capital at a coworking space."
    result = r3.score(script, "tier2_hindi_belt")
    assert len(result.invalid_signals_found) > 0


# ─── Step 3: R4 DebateResolutionReward ───────────────────────────────────────

def _make_critique_output(claims):
    from viral_script_engine.agents.critic import CritiqueOutput
    return CritiqueOutput(claims=claims, overall_severity="medium", raw_response="")


def _make_claim(claim_id, critique_class, timestamp_range, severity="high"):
    return CritiqueClaim(
        claim_id=claim_id,
        critique_class=critique_class,
        claim_text="test claim",
        timestamp_range=timestamp_range,
        evidence="evidence text",
        is_falsifiable=True,
        severity=severity,
    )


def _make_action():
    return ArbitratorAction(
        action_type=ActionType.HOOK_REWRITE,
        target_section="hook",
        instruction="fix hook",
        critique_claim_id="C1",
        reasoning="test",
    )


def test_r4_resolved_when_no_matching_claim():
    mock_critic = MagicMock()
    mock_critic.critique.return_value = _make_critique_output([
        _make_claim("C1", "cta_buried", "0:45-0:50"),
    ])
    r4 = DebateResolutionReward(critic_agent=mock_critic)
    original_claim = _make_claim("C1", "hook_weakness", "0:00-0:03", "high")
    result = r4.score("new script", _make_action(), original_claim,
                      "mumbai_gen_z", "instagram", "finance")
    assert result.score == 1.0
    assert result.resolution_status == "resolved"


def test_r4_partially_resolved_when_severity_drops():
    mock_critic = MagicMock()
    mock_critic.critique.return_value = _make_critique_output([
        _make_claim("C1", "hook_weakness", "0:01-0:04", "low"),
    ])
    r4 = DebateResolutionReward(critic_agent=mock_critic)
    original_claim = _make_claim("C1", "hook_weakness", "0:00-0:03", "high")
    result = r4.score("new script", _make_action(), original_claim,
                      "mumbai_gen_z", "instagram", "finance")
    assert result.score == 0.5
    assert result.resolution_status == "partially_resolved"


def test_r4_persists_when_same_severity():
    mock_critic = MagicMock()
    mock_critic.critique.return_value = _make_critique_output([
        _make_claim("C1", "hook_weakness", "0:01-0:03", "high"),
    ])
    r4 = DebateResolutionReward(critic_agent=mock_critic)
    original_claim = _make_claim("C1", "hook_weakness", "0:00-0:03", "high")
    result = r4.score("new script", _make_action(), original_claim,
                      "mumbai_gen_z", "instagram", "finance")
    assert result.score == 0.0
    assert result.resolution_status == "persists"


# ─── Step 4: R5 DefenderPreservationReward ───────────────────────────────────

def _make_defender_output(quote: str) -> DefenderOutput:
    return DefenderOutput(
        core_strength="Strong opening",
        core_strength_quote=quote,
        defense_argument="Should be preserved.",
        flagged_critic_claims=[],
        regional_voice_elements=[],
    )


def test_r5_high_score_when_quote_present():
    r5 = DefenderPreservationReward()
    quote = "Let me tell you a secret about saving money every month."
    script = "Let me tell you a secret about saving money every month. This is the key insight."
    defender_out = _make_defender_output(quote)
    result = r5.score(defender_out, script)
    assert result.score >= 0.85


def test_r5_zero_score_when_quote_absent():
    r5 = DefenderPreservationReward()
    quote = "Completely different text that shares nothing with rewrite."
    script = "Today we discuss quantum physics and neutron stars in distant galaxies."
    defender_out = _make_defender_output(quote)
    result = r5.score(defender_out, script)
    assert result.score < 0.65


# ─── Step 5/6: AntiGamingLog and RewardAggregator ────────────────────────────

def _make_components(**kwargs) -> RewardComponents:
    defaults = dict(
        r1_hook_strength=0.7, r2_coherence=0.7,
        r3_cultural_alignment=0.7,
        r4_debate_resolution=None,
        r5_defender_preservation=None,
    )
    defaults.update(kwargs)
    rc = RewardComponents(**defaults)
    rc.compute_total()
    return rc


def test_anti_gaming_catastrophic_drop_zeroes_reward():
    aggregator = RewardAggregator()
    start = _make_components(r2_coherence=0.8)
    current = _make_components(r2_coherence=0.4)

    result, log = aggregator.compute(current, start, [], episode_id="ep1", step_num=1)

    assert result.total == 0.0
    assert log.triggered is True
    assert log.rule_triggered == "catastrophic_drop"
    assert log.component_that_dropped == "r2_coherence"
    assert log.post_penalty_total == 0.0


def test_anti_gaming_diversity_penalty_fires_on_3x_same():
    aggregator = RewardAggregator()
    start = _make_components()
    current = _make_components()
    history = [ActionType.HOOK_REWRITE] * 3

    result, log = aggregator.compute(current, start, history, episode_id="ep2", step_num=2)

    assert log.triggered is True
    assert log.rule_triggered == "action_repetition"
    assert log.penalty_applied == 0.15


def test_anti_gaming_log_not_triggered_clean():
    aggregator = RewardAggregator()
    start = _make_components()
    current = _make_components()
    history = [ActionType.HOOK_REWRITE, ActionType.CTA_PLACEMENT, ActionType.SECTION_REORDER]

    result, log = aggregator.compute(current, start, history, episode_id="ep3", step_num=1)

    assert log.triggered is False
    assert log.rule_triggered is None
    assert log.penalty_applied == 0.0


def test_anti_gaming_log_fields_populated():
    aggregator = RewardAggregator()
    start = _make_components(r1_hook_strength=0.9)
    current = _make_components(r1_hook_strength=0.5)

    _, log = aggregator.compute(current, start, [], episode_id="myep", step_num=3)

    assert log.episode_id == "myep"
    assert log.step_num == 3
    assert isinstance(log, AntiGamingLog)