Blame for environment.py

"""
environment.py
==============
OpenEnv-compliant InformationIntegrityEnv.

The environment is 100% offline and deterministic once a scenario has been
collected by data_collector.py.  It accepts any scenario dict produced by
DataCollector.collect() — no hardcoded content.

OpenEnv contract
----------------
  env = InformationIntegrityEnv(scenario)
  obs               = env.reset()
  obs, rew, done, info = env.step(action)
  state_dict        = env.state()

Action space
------------
  verify_claim  claim_id                  → VerificationResult
  check_source  source_id                 → SourceEvidence
  decide        decision (allow/flag/remove) → episode ends

Observation
-----------
  content, claims (status masked), available_sources,
  interaction_history, status_flags, step_count

Design invariants
-----------------
  - Ground truth NEVER exposed in Observation
  - All randomness eliminated (same action sequence → same result)
  - Reward computed by reward.py (configurable, not hardcoded here)
  - Scenario loaded once at __init__; reset() clears episode state only
"""

from __future__ import annotations

import copy
import os
import time
import uuid
from typing import Any, Dict, List, Optional, Tuple

import yaml

from models import (
    Action,
    ActionType,
    Claim,
    DecisionType,
    EpisodeResult,
    InteractionRecord,
    Observation,
    Reward,
    RewardBreakdown,
    Scenario,
    ScenarioClaim,
    ScenarioSource,
    SourceInfo,
    Stance,
    StatusFlags,
    VerificationStatus,
)
from reward import EpisodeRewardState, RewardFunction

try:
    from retrieval import RetrievalPipeline
    from reasoning import DeterministicReasoningEngine, reason_over_claims
    _RETRIEVAL_AVAILABLE = True
except Exception:
    _RETRIEVAL_AVAILABLE = False

RELIABILITY_CONFIG_PATH = os.getenv("RELIABILITY_CONFIG", "domain_reliability.yaml")


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

def _load_evidence_thresholds() -> dict:
    if os.path.exists(RELIABILITY_CONFIG_PATH):
        with open(RELIABILITY_CONFIG_PATH, "r", encoding="utf-8") as fh:
            cfg = yaml.safe_load(fh) or {}
        return cfg.get("evidence_thresholds", {})
    return {}


def _scenario_from_dict(d: dict) -> Scenario:
    """Deserialise a raw scenario dict (e.g. from DataCollector) into a Scenario model."""
    claims = [ScenarioClaim(**c) for c in d["claims"]]
    sources = [ScenarioSource(**s) for s in d["sources"]]
    return Scenario(
        id=d.get("id", str(uuid.uuid4())),
        difficulty=d.get("difficulty", "medium"),
        content=d["content"],
        claims=claims,
        sources=sources,
        ground_truth=d["ground_truth"],
        topic=d.get("topic", ""),
        tags=d.get("tags", []),
        explanation_hint=d.get("explanation_hint", ""),
    )


# ---------------------------------------------------------------------------
# Main environment class
# ---------------------------------------------------------------------------

class InformationIntegrityEnv:
    """
    OpenEnv environment for reasoning-based content moderation.

    Parameters
    ----------
    scenario : dict | Scenario
        A collected scenario (from DataCollector) or a pre-built Scenario object.
    max_steps : int
        Maximum steps per episode (default 20).
    reward_fn : RewardFunction | None
        Pluggable reward function (default: RewardFunction()).
    """

    # Public metadata (read by openenv validate)
    ENV_NAME    = "InformationIntegrityEnv"
    ENV_VERSION = "1.0.0"
    ENV_AUTHOR  = "openenv-content-moderation"

    def __init__(
        self,
        scenario: dict | Scenario,
        max_steps: int = 20,
        reward_fn: Optional[RewardFunction] = None,
    ) -> None:
        if isinstance(scenario, dict):
            scenario = _scenario_from_dict(scenario)
        self._scenario   : Scenario          = scenario
        self._max_steps  : int               = max_steps
        self._reward_fn  : RewardFunction    = reward_fn or RewardFunction()
        self._thresholds : dict              = _load_evidence_thresholds()

        # Episode state (initialised by reset)
        self._episode_id : str               = ""
        self._claims     : List[Claim]       = []
        self._history    : List[InteractionRecord] = []
        self._flags      : StatusFlags       = StatusFlags()
        self._step_count : int               = 0
        self._done       : bool              = False
        self._final_decision: Optional[str]  = None
        self._reward_state: EpisodeRewardState = EpisodeRewardState()

        # Index for fast lookup
        self._source_index : Dict[str, ScenarioSource] = {
            s.id: s for s in self._scenario.sources
        }
        self._claim_index  : Dict[str, ScenarioClaim]  = {
            c.id: c for c in self._scenario.claims
        }

        # Retrieval + reasoning engines (lazy init for pickle compat)
        self._retrieval_pipeline: Optional[Any] = None
        self._reasoning_engine:   Optional[Any] = None
        self._last_reasoning:     Optional[Any] = None  # StructuredReasoning cached

    # ------------------------------------------------------------------
    # OpenEnv API
    # ------------------------------------------------------------------

    def reset(
        self,
        difficulty: Optional[str] = None,
        scenario_id: Optional[str | int] = None,
    ) -> Observation:
        """
        Reset episode state and return initial observation.

        Programmatic difficulty control (no UI dependency):

        Parameters
        ----------
        difficulty : str | None
            If provided, load a new scenario from the deterministic dataset
            for that difficulty tier ('easy', 'medium', 'hard').
            Overrides the scenario supplied at __init__ time.
        scenario_id : str | int | None
            Index (int) or scenario id (str) within the difficulty dataset.
            If None and difficulty is provided, a random scenario is selected.
            Ignored when difficulty is None.
        """
        if difficulty is not None:
            from tasks.registry import TASKS
            import random as _random
            task = TASKS[difficulty]
            if scenario_id is None:
                sc_dict = task.sample()
            else:
                sc_dict = task.get_scenario(scenario_id)
            self._scenario    = _scenario_from_dict(sc_dict)
            self._source_index = {s.id: s for s in self._scenario.sources}
            self._claim_index  = {c.id: c for c in self._scenario.claims}

        self._episode_id      = str(uuid.uuid4())
        self._step_count      = 0
        self._done            = False
        self._final_decision  = None
        self._history         = []
        self._flags           = StatusFlags()
        self._reward_state    = EpisodeRewardState()
        self._last_reasoning  = None

        # Reinitialise mutable claim states (no ground-truth exposed)
        self._claims = [
            Claim(id=c.id, text=c.text)
            for c in self._scenario.claims
        ]

        return self._build_observation()

    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
        """
        Process one agent action.

        Returns
        -------
        observation : Observation
        reward      : Reward
        done        : bool
        info        : dict  (extra diagnostics — not for agent use)
        """
        if self._done:
            raise RuntimeError(
                "Episode has ended. Call reset() to start a new episode."
            )

        self._step_count += 1
        self._reward_state.step_count = self._step_count

        if action.type == ActionType.VERIFY_CLAIM:
            step_reward, result = self._handle_verify_claim(action)
        elif action.type == ActionType.CHECK_SOURCE:
            step_reward, result = self._handle_check_source(action)
        elif action.type == ActionType.DECIDE:
            step_reward, result = self._handle_decide(action)
        else:
            raise ValueError(f"Unknown action type: {action.type}")

        # Max-steps check
        if self._step_count >= self._max_steps and not self._done:
            self._done = True
            self._flags.max_steps_reached = True
            result["max_steps_reached"] = True

        # Update evidence sufficiency flag
        self._flags.sufficient_evidence = self._check_sufficient_evidence()

        # Log interaction
        record = InteractionRecord(
            step=self._step_count,
            action_type=action.type,
            action_payload=action.model_dump(exclude_none=True),
            result=result,
            reward_delta=step_reward,
        )
        self._history.append(record)

        obs    = self._build_observation()
        reward = self._build_reward(step_reward, result.get("message", ""))
        info   = {
            "episode_id": self._episode_id,
            "step":       self._step_count,
            "done":       self._done,
            **result,
        }
        return obs, reward, self._done, info

    def state(self) -> Dict[str, Any]:
        """Return full internal state (for debugging / logging — not for agent)."""
        return {
            "episode_id":       self._episode_id,
            "scenario_id":      self._scenario.id,
            "difficulty":       self._scenario.difficulty,
            "step_count":       self._step_count,
            "done":             self._done,
            "final_decision":   self._final_decision,
            "ground_truth":     self._scenario.ground_truth,  # hidden from agent obs
            "claims": [
                {
                    "id":                  c.id,
                    "text":                c.text,
                    "verification_status": c.verification_status,
                    "confidence":          c.confidence,
                }
                for c in self._claims
            ],
            "reward_breakdown": self._reward_fn.breakdown_dict(self._reward_state),
            "flags":            self._flags.model_dump(),
            "history_len":      len(self._history),
            "last_reasoning":   (
                {ct: r.to_dict() for ct, r in self._last_reasoning.items()}
                if self._last_reasoning else None
            ),
        }

    def episode_result(self, explanation: Optional[str] = None) -> EpisodeResult:
        """Build a final EpisodeResult (call after done=True)."""
        bd = self._reward_fn.breakdown_dict(self._reward_state)
        ground_truth = self._scenario.ground_truth
        correct = self._final_decision == ground_truth

        from tasks.graders import GraderFactory
        grader = GraderFactory.get(self._scenario.difficulty)
        grader_score = grader.score(
            decision=self._final_decision or "none",
            ground_truth=ground_truth,
            sources_checked=list(self._reward_state.checked_sources),
            claims_verified=list(self._reward_state.verified_claims),
            step_count=self._step_count,
            all_source_ids=[s.id for s in self._scenario.sources],
            all_claim_ids=[c.id for c in self._scenario.claims],
        )

        return EpisodeResult(
            episode_id=self._episode_id,
            task_difficulty=self._scenario.difficulty,
            scenario_index=0,
            total_reward=self._reward_fn.final_score(self._reward_state),
            decision=self._final_decision,
            ground_truth=ground_truth,
            correct=correct,
            steps_taken=self._step_count,
            claims_verified=len(self._reward_state.verified_claims),
            sources_checked=len(self._reward_state.checked_sources),
            grader_score=grader_score,
            reward_breakdown=RewardBreakdown(
                verify_claim_reward=bd["verify_claim_reward"],
                check_source_reward=bd["check_source_reward"],
                multi_source_bonus=bd["multi_source_bonus"],
                decision_reward=bd["decision_reward"],
                premature_penalty=bd["premature_penalty"],
                redundant_penalty=bd["redundant_penalty"],
                inefficiency_penalty=bd["inefficiency_penalty"],
                aggression_penalty=bd["aggression_penalty"],
            ),
            explanation=explanation,
        )

    # ------------------------------------------------------------------
    # Action handlers
    # ------------------------------------------------------------------

    def _handle_verify_claim(self, action: Action) -> Tuple[float, dict]:
        claim_id = action.claim_id
        if claim_id not in self._claim_index:
            return 0.0, {"error": f"Unknown claim_id: {claim_id}"}

        scenario_claim = self._claim_index[claim_id]
        agent_claim    = self._get_agent_claim(claim_id)

        # Reveal verification status from scenario ground truth (masked from Observation)
        status     = scenario_claim.ground_truth_status
        confidence = scenario_claim.ground_truth_confidence

        # Use RetrievalPipeline to find the most relevant sources for THIS claim
        if _RETRIEVAL_AVAILABLE:
            if self._retrieval_pipeline is None:
                self._retrieval_pipeline = RetrievalPipeline()
            raw_sources = [
                {
                    "id":          s.id,
                    "title":       s.title,
                    "text":        s.text,
                    "url":         s.url,
                    "stance":      s.stance,
                    "reliability": s.reliability,
                    "source_type": s.source_type,
                }
                for s in self._scenario.sources
            ]
            ranked = self._retrieval_pipeline.retrieve(scenario_claim.text, raw_sources)
            relevant_sources = [s.id for s in ranked[:5]]
        else:
            relevant_sources = [
                s.id for s in self._scenario.sources
                if s.stance in ("contradicting", "supporting")
            ]

        # Update mutable claim object
        agent_claim.verification_status = VerificationStatus(status)
        agent_claim.confidence          = confidence
        agent_claim.verified_by_sources = relevant_sources[:3]

        # Reward
        r, msg = self._reward_fn.compute_verify_claim(claim_id, self._reward_state)

        result: dict = {
            "claim_id":          claim_id,
            "status":            status,
            "confidence":        confidence,
            "message":           msg,
            "relevant_sources":  relevant_sources,
            "hint":              f"Claim '{scenario_claim.text[:60]}…' assessed as {status}",
        }
        return r, result

    def _handle_check_source(self, action: Action) -> Tuple[float, dict]:
        source_id = action.source_id
        if source_id not in self._source_index:
            return 0.0, {"error": f"Unknown source_id: {source_id}"}

        src = self._source_index[source_id]

        # Reward
        r, msg = self._reward_fn.compute_check_source(source_id, self._reward_state)

        result: dict = {
            "source_id":   source_id,
            "title":       src.title,
            "text":        src.text,
            "url":         src.url,
            "stance":      src.stance,
            "reliability": src.reliability,
            "source_type": src.source_type,
            "message":     msg,
        }
        return r, result

    def _handle_decide(self, action: Action) -> Tuple[float, dict]:
        decision     = action.decision
        ground_truth = self._scenario.ground_truth
        num_claims   = len(self._scenario.claims)

        # Run deterministic reasoning engine if available
        reasoning_summary = {}
        if _RETRIEVAL_AVAILABLE:
            if self._retrieval_pipeline is None:
                self._retrieval_pipeline = RetrievalPipeline()
            if self._reasoning_engine is None:
                self._reasoning_engine = DeterministicReasoningEngine()

            raw_sources = [
                {
                    "id":          s.id,
                    "title":       s.title,
                    "text":        s.text,
                    "url":         s.url,
                    "stance":      s.stance,
                    "reliability": s.reliability,
                    "source_type": s.source_type,
                }
                for s in self._scenario.sources
            ]
            claim_texts = [c.text for c in self._scenario.claims]
            claim_results = self._retrieval_pipeline.retrieve_for_all_claims(
                claim_texts, raw_sources
            )
            agg_decision, per_claim_reasoning = reason_over_claims(
                claim_results, self._reasoning_engine
            )
            # Cache reasoning for episode_result()
            self._last_reasoning = per_claim_reasoning
            reasoning_summary = {
                "engine_decision": agg_decision,
                "agent_decision":  decision,
                "per_claim": {
                    ct: r.to_dict() for ct, r in per_claim_reasoning.items()
                },
            }

        r, msg = self._reward_fn.compute_decision(
            decision=decision,
            ground_truth=ground_truth,
            state=self._reward_state,
            num_claims=num_claims,
        )

        self._final_decision      = decision
        self._done                = True
        self._flags.decision_made = True

        if "premature" in msg:
            self._flags.premature_decision = True

        result: dict = {
            "decision":        decision,
            "message":         msg,
            "episode_done":    True,
            "reasoning":       reasoning_summary,
        }
        return r, result

    # ------------------------------------------------------------------
    # State construction helpers
    # ------------------------------------------------------------------

    def _build_observation(self) -> Observation:
        """Build agent-facing Observation — ground truth is NEVER included."""
        sources_view = [
            SourceInfo(
                id=s.id,
                title=s.title,
                text=s.text,
                url=s.url,
                stance=Stance(s.stance),
                reliability=s.reliability,
                source_type=s.source_type,
            )
            for s in self._scenario.sources
        ]
        return Observation(
            content=self._scenario.content,
            claims=copy.deepcopy(self._claims),
            available_sources=sources_view,
            interaction_history=list(self._history),
            status_flags=copy.deepcopy(self._flags),
            step_count=self._step_count,
            episode_id=self._episode_id,
            task_difficulty=self._scenario.difficulty,
        )

    def _build_reward(self, score: float, message: str) -> Reward:
        bd = self._reward_fn.breakdown_dict(self._reward_state)
        return Reward(
            score=score,
            breakdown=RewardBreakdown(
                verify_claim_reward=bd["verify_claim_reward"],
                check_source_reward=bd["check_source_reward"],
                multi_source_bonus=bd["multi_source_bonus"],
                decision_reward=bd["decision_reward"],
                premature_penalty=bd["premature_penalty"],
                redundant_penalty=bd["redundant_penalty"],
                inefficiency_penalty=bd["inefficiency_penalty"],
                aggression_penalty=bd["aggression_penalty"],
            ),
            message=message,
            is_terminal=self._done,
        )

    def _get_agent_claim(self, claim_id: str) -> Claim:
        for c in self._claims:
            if c.id == claim_id:
                return c
        raise KeyError(f"Claim {claim_id} not found in agent state")

    def _check_sufficient_evidence(self) -> bool:
        thresh = self._thresholds
        min_sources  = int(thresh.get("min_sources_for_decision", 2))
        min_verified = float(thresh.get("min_claims_verified_fraction", 0.5))
        n_claims     = len(self._scenario.claims)
        n_verified   = len(self._reward_state.verified_claims)
        n_sources    = len(self._reward_state.checked_sources)

        frac_ok = (n_verified / n_claims >= min_verified) if n_claims > 0 else True
        return n_sources >= min_sources and frac_ok

    # ------------------------------------------------------------------
    # Convenience helpers
    # ------------------------------------------------------------------

    @property
    def scenario(self) -> Scenario:
        return self._scenario

    @property
    def max_steps(self) -> int:
        return self._max_steps

    @property
    def done(self) -> bool:
        return self._done

    @classmethod
    def from_dataset(
        cls,
        difficulty: str,
        scenario_id: Optional[str | int] = None,
        max_steps: int = 20,
        reward_fn: Optional[RewardFunction] = None,
    ) -> "InformationIntegrityEnv":
        """
        Build an env directly from the deterministic dataset — no DataCollector needed.

        Parameters
        ----------
        difficulty : 'easy' | 'medium' | 'hard'
        scenario_id : int index or str id; None = random selection
        max_steps   : episode step limit
        reward_fn   : optional custom RewardFunction

        Example
        -------
        env = InformationIntegrityEnv.from_dataset("medium", scenario_id=2)
        obs = env.reset()
        """
        from tasks.registry import TASKS
        task = TASKS[difficulty]
        if scenario_id is None:
            sc_dict = task.sample()
        else:
            sc_dict = task.get_scenario(scenario_id)
        return cls(scenario=sc_dict, max_steps=max_steps, reward_fn=reward_fn)

    @classmethod
    def from_collector(
        cls,
        content: str,
        claims: List[str],
        difficulty: str = "medium",
        **collector_kwargs,
    ) -> "InformationIntegrityEnv":
        """
        Convenience factory: collect a scenario on-the-fly and build an env.

        Example
        -------
        env = InformationIntegrityEnv.from_collector(
            content="Post text here",
            claims=["Claim A", "Claim B"],
            difficulty="hard",
        )
        obs = env.reset()
        """
        from data_collector import CollectionRequest, DataCollector

        collector = DataCollector()
        request   = CollectionRequest(
            content=content,
            claims=claims,
            difficulty=difficulty,
            **collector_kwargs,
        )
        scenario_dict = collector.collect(request)
        return cls(scenario=scenario_dict)