viral_script_engine/environment/env.py

import json
import random
import time
from collections import Counter
from typing import Optional, Tuple

from viral_script_engine.agents.critic import CriticAgent
from viral_script_engine.agents.defender import DefenderAgent
from viral_script_engine.agents.rewriter import RewriterAgent
from viral_script_engine.agents.reasoning_parser import ReasoningParser, ArbitratorParseError
from viral_script_engine.environment.actions import ArbitratorAction
from viral_script_engine.environment.episode_state import EpisodeState
from viral_script_engine.environment.observations import (
    DebateRound, Observation, RewardComponents,
)
from viral_script_engine.agents.moderation_agent import ModerationAgent
from viral_script_engine.agents.originality_agent import OriginalityAgent
from viral_script_engine.rewards.r1_hook_strength import HookStrengthReward
from viral_script_engine.rewards.r2_coherence import CoherenceReward
from viral_script_engine.rewards.r3_cultural_alignment import CulturalAlignmentReward
from viral_script_engine.rewards.r4_debate_resolution import DebateResolutionReward
from viral_script_engine.rewards.r5_defender_preservation import DefenderPreservationReward
from viral_script_engine.rewards.r6_safety import SafetyReward
from viral_script_engine.rewards.r7_originality import OriginalityReward
from viral_script_engine.rewards.reward_aggregator import RewardAggregator
from viral_script_engine.rewards.process_reward import ProcessReward, ProcessRewardResult
from viral_script_engine.personas.creator_profile import CreatorProfile, CreatorTier
from viral_script_engine.personas.profile_generator import ProfileGenerator
from viral_script_engine.rewards.r8_persona_fit import PersonaFitReward
from viral_script_engine.rewards.r9_platform_pacing import PlatformPacingReward
from viral_script_engine.platforms.platform_spec import PlatformRegistry
from viral_script_engine.memory.memory_compressor import MemoryCompressor
from viral_script_engine.memory.history_store import HistoryStore
from viral_script_engine.rewards.r10_retention_curve import RetentionCurveReward

_TIERS = {
    "easy": ["S01", "S02", "S03", "S04"],
    "medium": ["S05", "S06", "S07"],
    "hard": ["S08", "S09", "S10"],
    "self_generated": [],
}


class ViralScriptEnv:
    def __init__(
        self,
        scripts_path: str = "data/test_scripts/scripts.json",
        max_steps: int = 5,
        difficulty: str = "easy",
        use_anti_gaming: bool = True,
        cultural_kb_path: str = "data/cultural_kb.json",
        use_escalation: bool = True,
        difficulty_tracker=None,
        escalation_engine=None,
    ):
        self.max_steps = max_steps
        self.difficulty = difficulty
        self.use_anti_gaming = use_anti_gaming
        self.use_escalation = use_escalation

        with open(scripts_path) as f:
            all_scripts = json.load(f)

        tier_ids = _TIERS.get(difficulty, [])
        self._scripts = [s for s in all_scripts if s["script_id"] in tier_ids]
        if not self._scripts:
            self._scripts = all_scripts

        self.critic = CriticAgent(backend="hf")
        self.defender = DefenderAgent(backend="hf")
        self.rewriter = RewriterAgent(backend="hf")
        self.r1 = HookStrengthReward()
        self.r2 = CoherenceReward()
        self.r3 = CulturalAlignmentReward(knowledge_base_path=cultural_kb_path)
        self.r4 = DebateResolutionReward(critic_agent=self.critic)
        self.r5 = DefenderPreservationReward()
        self.r6 = SafetyReward()
        self.r7 = OriginalityReward()
        self.moderation_agent = ModerationAgent()
        self.originality_agent = OriginalityAgent()
        self.aggregator = RewardAggregator()
        self.reasoning_parser = ReasoningParser()
        self.process_reward_calc = ProcessReward()
        self.profile_generator = ProfileGenerator()
        self.r8 = PersonaFitReward()
        self.r9 = PlatformPacingReward()
        self.platform_registry = PlatformRegistry()
        self.memory_compressor = MemoryCompressor()
        self.history_store = HistoryStore()
        self.r10 = RetentionCurveReward(cultural_kb_path=cultural_kb_path)
        self._state: Optional[EpisodeState] = None
        self._current_profile: Optional[CreatorProfile] = None
        self._current_platform: str = "Reels"
        self._current_creator_id: str = "default"
        self._current_history_buffer = None

        if use_escalation:
            if difficulty_tracker is None:
                from viral_script_engine.escalation.difficulty_tracker import DifficultyTracker
                difficulty_tracker = DifficultyTracker()
            if escalation_engine is None:
                from viral_script_engine.escalation.critic_escalation_engine import CriticEscalationEngine
                escalation_engine = CriticEscalationEngine()

        self.difficulty_tracker = difficulty_tracker
        self.escalation_engine = escalation_engine

        # Track first-step critic output per episode for dominant class detection
        self._first_critique = None
        self._timeout_count: int = 0

    def reset_from_config(self, episode_config: dict) -> Tuple[dict, dict]:
        """Reset the environment to a specific episode config from curriculum JSONL."""
        script = {
            "script_id": episode_config.get("script_id", "unknown"),
            "script_text": episode_config["script_text"],
            "region": episode_config["region"],
            "platform": episode_config["platform"],
            "niche": episode_config["niche"],
        }
        return self._reset_with_script(script, episode_config.get("difficulty", self.difficulty))

    def reset(self, seed=None, options=None) -> Tuple[dict, dict]:
        if seed is not None:
            random.seed(seed)

        self._first_critique = None
        used_escalation = False

        if self.use_escalation and self.difficulty_tracker and self.escalation_engine:
            mastered = self.difficulty_tracker.get_mastered_classes()
            if mastered:
                challenge = self.escalation_engine.get_next_challenge(self.difficulty_tracker)
                if challenge is None:
                    # Generate a new escalated challenge from the first mastered class
                    src_class = mastered[0]
                    example_script = random.choice(self._scripts)
                    challenge = self.escalation_engine.escalate(
                        mastered_class=src_class,
                        original_script_example=example_script["script_text"],
                        region=example_script.get("region", "pan_india_english"),
                        platform=example_script.get("platform", "Reels"),
                    )

                script = challenge.to_script_dict()
                print(f"[ESCALATION] Using self-generated challenge for class '{challenge.source_class}' — {challenge.why_its_harder}")
                obs, info = self._reset_with_script(script, "self_generated")
                info["escalation_used"] = True
                info["escalation_source_class"] = challenge.source_class
                return obs, info

        script = random.choice(self._scripts)
        obs, info = self._reset_with_script(script, self.difficulty)
        info["escalation_used"] = False
        return obs, info

    def _reset_with_script(self, script: dict, difficulty: str) -> Tuple[dict, dict]:
        self._current_creator_id = script.get("creator_id", script.get("script_id", "default"))
        self._current_history_buffer = self.history_store.load(self._current_creator_id)
        self._current_platform = script.get("platform", "Reels")
        r1_result = self.r1.score(script["script_text"], platform=self._current_platform)
        r2_result = self.r2.score(script["script_text"], script["script_text"], platform=self._current_platform)
        r3_result = self.r3.score(script["script_text"], script.get("region", "pan_india_english"))
        mod_out = self.moderation_agent.check(script["script_text"])
        orig_out = self.originality_agent.check(script["script_text"])
        r6_result = self.r6.score(mod_out)
        r7_result = self.r7.score(orig_out)
        initial_rewards = RewardComponents(
            r1_hook_strength=r1_result.score,
            r2_coherence=r2_result.score,
            r3_cultural_alignment=r3_result.score,
            r6_safety=r6_result.score,
            r7_originality=r7_result.score,
        )
        initial_rewards.compute_total()

        self._state = EpisodeState.new(
            script=script,
            max_steps=self.max_steps,
            difficulty_level=difficulty,
            initial_rewards=initial_rewards,
        )

        # Phase 8: generate a creator profile matching episode difficulty
        self._current_profile = self._generate_profile_for_difficulty(
            difficulty=difficulty,
            niche=script.get("niche", "personal finance"),
            seed=hash(script.get("script_id", "default")) % (2 ** 31),
        )

        return self._build_observation().model_dump(), {}

    def _generate_profile_for_difficulty(
        self, difficulty: str, niche: str, seed: int
    ) -> CreatorProfile:
        """Map episode difficulty to an appropriate creator tier."""
        tier_map = {
            "easy": [CreatorTier.BEGINNER, CreatorTier.GROWING],
            "medium": [CreatorTier.GROWING, CreatorTier.ESTABLISHED],
            "hard": [CreatorTier.ESTABLISHED, CreatorTier.VERIFIED],
            "self_generated": [CreatorTier.ESTABLISHED, CreatorTier.VERIFIED],
        }
        import random as _rng
        tiers = tier_map.get(difficulty, [CreatorTier.GROWING])
        tier = _rng.Random(seed).choice(tiers)
        return self.profile_generator.generate(tier=tier, niche=niche, seed=seed)

    def step(self, action: dict, raw_output: str = None) -> Tuple[dict, float, bool, bool, dict]:
        if self._state is None:
            raise RuntimeError("Call reset() before step()")

        _step_start = time.time()

        arb_action = ArbitratorAction(**action)

        try:
            critique = self.critic.critique(
                script=self._state.current_script,
                region=self._state.region,
                platform=self._state.platform,
                niche=self._state.niche,
            )
        except TimeoutError:
            self._timeout_count += 1
            info = {"timeout": True, "timeout_agent": "critic", "timeout_count": self._timeout_count}
            return self._build_observation().model_dump(), 0.0, False, True, info

        # Track first critique for dominant class detection at episode end
        if self._state.step_num == 0:
            self._first_critique = critique

        try:
            defender_output = self.defender.defend(
                script=self._state.current_script,
                critic_claims=critique.claims,
                region=self._state.region,
                platform=self._state.platform,
            )
        except TimeoutError:
            self._timeout_count += 1
            info = {"timeout": True, "timeout_agent": "defender", "timeout_count": self._timeout_count}
            return self._build_observation().model_dump(), 0.0, False, True, info

        # Phase 7: parse reasoning chain and compute process reward before rewrite
        reasoning_chain = None
        process_result = None
        if raw_output:
            try:
                reasoning_chain = self.reasoning_parser.parse(raw_output)
                process_result = self.process_reward_calc.score(
                    reasoning_chain=reasoning_chain,
                    critic_claims=critique.claims,
                    defender_output=defender_output,
                    current_reward_components=self._state.last_reward_components,
                    episode_start_components=self._state.episode_start_rewards,
                )
            except ArbitratorParseError:
                reasoning_chain = None
                process_result = None

        try:
            rewrite_result = self.rewriter.rewrite(self._state.current_script, arb_action)
        except TimeoutError:
            self._timeout_count += 1
            info = {"timeout": True, "timeout_agent": "rewriter", "timeout_count": self._timeout_count}
            return self._build_observation().model_dump(), 0.0, False, True, info
        new_script = rewrite_result.rewritten_script

        r1_result = self.r1.score(new_script, platform=self._current_platform)
        r2_result = self.r2.score(self._state.original_script, new_script, platform=self._current_platform)
        r3_result = self.r3.score(new_script, self._state.region)

        targeted_claim = next(
            (c for c in critique.claims if c.claim_id == arb_action.critique_claim_id),
            critique.claims[0] if critique.claims else None,
        )
        r4_result = self.r4.score(
            new_script=new_script,
            original_action=arb_action,
            original_claim=targeted_claim,
            region=self._state.region,
            platform=self._state.platform,
            niche=self._state.niche,
        ) if targeted_claim else None

        r5_result = self.r5.score(defender_output, new_script)

        moderation_out = self.moderation_agent.check(new_script)
        originality_out = self.originality_agent.check(new_script)
        r6_result = self.r6.score(moderation_out)
        r7_result = self.r7.score(originality_out)

        # Phase 8: compute R8 persona fit
        r8_score = None
        if self._current_profile is not None and targeted_claim is not None:
            r8_result = self.r8.score(
                action=arb_action,
                creator_profile=self._current_profile,
                addressed_critique_class=targeted_claim.critique_class,
            )
            r8_score = r8_result.score

        # Phase 9: compute R9 platform pacing
        r9_result = self.r9.score(new_script, platform=self._current_platform)

        # Phase 12: compute R10 retention curve reward
        r10_score = None
        if self.r10.predictor._trained:
            try:
                r10_result = self.r10.score(
                    original_script=self._state.original_script,
                    rewritten_script=new_script,
                    platform=self._current_platform,
                    region=self._state.region,
                    action_type=str(arb_action.action_type.value),
                    episode_id=self._state.episode_id,
                )
                r10_score = r10_result.score
            except Exception:
                r10_score = None

        components = RewardComponents(
            r1_hook_strength=r1_result.score,
            r2_coherence=r2_result.score,
            r3_cultural_alignment=r3_result.score,
            r4_debate_resolution=r4_result.score if r4_result else None,
            r5_defender_preservation=r5_result.score,
            r6_safety=r6_result.score,
            r7_originality=r7_result.score,
            r8_persona_fit=r8_score,
            r9_platform_pacing=r9_result.score,
            r10_retention_curve=r10_score,
            process_reward=process_result.weighted_contribution if process_result else None,
        )

        self._state.action_history.append(arb_action.action_type)
        if self.use_anti_gaming:
            components, anti_log = self.aggregator.compute(
                components,
                self._state.episode_start_rewards,
                self._state.action_history,
                episode_id=self._state.episode_id,
                step_num=self._state.step_num,
            )
        else:
            components.compute_total()
            from viral_script_engine.rewards.reward_aggregator import AntiGamingLog
            anti_log = AntiGamingLog(
                episode_id=self._state.episode_id,
                step_num=self._state.step_num,
                triggered=False,
                penalty_applied=0.0,
                pre_penalty_total=components.total,
                post_penalty_total=components.total,
            )

        round_ = DebateRound(
            step_num=self._state.step_num,
            critic_claims=critique.claims,
            defender_response=defender_output.model_dump(),
            arbitrator_action=arb_action,
            rewrite_diff=rewrite_result.diff,
            reward_components=components,
            moderation_output=moderation_out.model_dump(),
            originality_output=originality_out.model_dump(),
            reasoning_chain=reasoning_chain.model_dump() if reasoning_chain else None,
        )
        self._state.debate_history.append(round_)
        self._state.current_script = new_script
        self._state.last_reward_components = components
        self._state.step_num += 1

        if not hasattr(self._state, "anti_gaming_logs"):
            self._state.anti_gaming_logs = []
        self._state.anti_gaming_logs.append(anti_log.model_dump())

        terminated = (
            self._state.step_num >= self._state.max_steps
            or components.total >= 0.9
        )

        if terminated and self.use_escalation and self.difficulty_tracker:
            dominant_class = self._get_dominant_critique_class()
            r4_score = components.r4_debate_resolution if components.r4_debate_resolution is not None else 0.0
            self.difficulty_tracker.record_episode(
                dominant_critique_class=dominant_class,
                r4_score=r4_score,
                episode_id=self._state.episode_id,
            )

        if terminated:
            episode_number = (
                (self._current_history_buffer.total_episodes + 1)
                if self._current_history_buffer else 1
            )
            new_memory = self.memory_compressor.compress(
                episode_log=self._build_episode_log(),
                episode_number=episode_number,
            )
            self._current_history_buffer = self.memory_compressor.update_buffer(
                self._current_history_buffer, new_memory, self._current_creator_id
            )
            self.history_store.save(self._current_history_buffer)

        if time.time() - _step_start > 120:
            self._timeout_count += 1
            return self._build_observation().model_dump(), 0.0, False, True, {
                "timeout": True, "timeout_agent": "step_wall_clock",
                "timeout_count": self._timeout_count,
            }

        info = {
            "reward_components": components.model_dump(),
            "anti_gaming_triggered": anti_log.triggered,
            "penalty_reason": anti_log.rule_triggered,
            "anti_gaming_log": anti_log.model_dump(),
            "moderation_output": moderation_out.model_dump(),
            "originality_output": originality_out.model_dump(),
            "process_reward_result": process_result.model_dump() if process_result else None,
            "reasoning_chain": reasoning_chain.model_dump() if reasoning_chain else None,
            "creator_profile": self._current_profile.model_dump(mode="json") if self._current_profile else None,
            "timeout_count": self._timeout_count,
        }
        return self._build_observation().model_dump(), components.total, terminated, False, info

    def _build_episode_log(self) -> dict:
        s = self._state
        first_claims = []
        if self._first_critique and self._first_critique.claims:
            first_claims = [c.model_dump() for c in self._first_critique.claims]
        return {
            "episode_id": s.episode_id,
            "niche": s.niche,
            "platform": s.platform,
            "actions_taken": [a.value if hasattr(a, "value") else str(a) for a in s.action_history],
            "first_critique_claims": first_claims,
            "initial_reward_components": s.episode_start_rewards.model_dump(),
            "final_reward_components": s.last_reward_components.model_dump(),
            "final_total_reward": s.last_reward_components.total,
        }

    def _get_dominant_critique_class(self) -> str:
        """Return the most common critique_class from the first episode critique."""
        if self._first_critique is None or not self._first_critique.claims:
            return "hook_weakness"
        counts = Counter(c.critique_class for c in self._first_critique.claims)
        return counts.most_common(1)[0][0]

    def state(self) -> dict:
        if self._state is None:
            return {}
        s = self._state
        return {
            "current_script": s.current_script,
            "original_script": s.original_script,
            "debate_history": [r.model_dump() for r in s.debate_history],
            "reward_components": s.last_reward_components.model_dump(),
            "step_num": s.step_num,
            "difficulty_level": s.difficulty_level,
            "episode_id": s.episode_id,
            "anti_gaming_logs": getattr(s, "anti_gaming_logs", []),
            "creator_profile": self._current_profile.model_dump(mode="json") if self._current_profile else None,
            "timeout_count": self._timeout_count,
        }

    def _build_observation(self) -> Observation:
        s = self._state
        last_round = s.debate_history[-1] if s.debate_history else None
        mod_flags = []
        orig_flags = []
        if last_round and last_round.moderation_output:
            mod_flags = last_round.moderation_output.get("flags", [])
        if last_round and last_round.originality_output:
            orig_flags = last_round.originality_output.get("flags", [])
        history_context = (
            self._current_history_buffer.to_prompt_context()
            if self._current_history_buffer else None
        )
        return Observation(
            current_script=s.current_script,
            original_script=s.original_script,
            region=s.region,
            platform=s.platform,
            niche=s.niche,
            step_num=s.step_num,
            max_steps=s.max_steps,
            debate_history=s.debate_history,
            reward_components=s.last_reward_components,
            difficulty_level=s.difficulty_level,
            episode_id=s.episode_id,
            current_moderation_flags=mod_flags,
            current_originality_flags=orig_flags,
            creator_profile=self._current_profile.model_dump(mode="json") if self._current_profile else None,
            creator_history=self._current_history_buffer.model_dump() if self._current_history_buffer else None,
            history_context=history_context,
        )