server/talking_head_environment.py

"""
TalkingHeadBench OpenEnv Environment.

Wraps the existing benchmark into a classic OpenEnv step-based environment.
Episode flow is intentionally simplified to four agent-facing steps:

1) reset() emits ImageDiagnosticsObservation (Node 1 context)
2) step(ImageDiagnosticsAction) emits ParamAnomalyObservation (Node 2 context)
3) step(ParamAnomalyAction) emits PhonemeRiskObservation (Node 8 context)
4) step(PhonemeRiskAction) returns done=True with final weighted reward

Final reward:
    final = 0.25 * subenv1 + 0.35 * subenv2 + 0.40 * subenv3
"""

from __future__ import annotations

import json
import logging
import random
from pathlib import Path
from typing import Any, Optional, TypeVar
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import Action, Observation, State

from src.envs.subenv1.node2_param_anomaly import detect_param_anomalies
from src.envs.subenv1.node3_grader import (
    grade_anomaly_detection,
    grade_image_diagnostics,
    produce_reference_audit_handoff,
)
from src.envs.subenv2.node5_disposition import recommend_clip_disposition
from src.envs.subenv2.node6_grader import grade_clip_disposition
from src.envs.subenv3.node9_grader import grade_behavioral_audit
from src.schemas.ground_truth import (
    GroundTruthBehavioralAnnotation,
    GroundTruthClipAnnotation,
    GroundTruthImageAnnotation,
    GroundTruthParamAnnotation,
)
from src.schemas.subenv1 import (
    ImageDiagnosticsAction,
    ImageDiagnosticsObservation,
    ParamAnomalyAction,
    ParamAnomalyObservation,
)
from src.schemas.subenv2 import (
    ClipDispositionObservation,
    ClipEvidenceDossier,
    ClipSignalObservation,
)
from src.schemas.subenv3 import PhonemeRiskAction, PhonemeRiskObservation

log = logging.getLogger(__name__)

_T = TypeVar("_T")
_TEST_SET_DIR = Path(__file__).resolve().parent.parent / "tests" / "test_set"


class TalkingHeadEnvironment(Environment[Action, Observation, State]):
    """OpenEnv adapter for TalkingHeadBench."""

    def __init__(self) -> None:
        self._state = State(episode_id=str(uuid4()), step_count=0)

        self._subenv1_cases = self._load_cases("subenv1_cases.json")
        self._subenv2_cases = self._load_cases("subenv2_cases.json")
        self._subenv3_cases = self._load_cases("subenv3_cases.json")

        self._cursor = 0
        self._phase = "idle"
        self._scores: dict[str, float] = {}

        self._current_case: dict[str, Any] | None = None
        self._node1_obs: ImageDiagnosticsObservation | None = None
        self._node1_gt: GroundTruthImageAnnotation | None = None
        self._node1_action: ImageDiagnosticsAction | None = None

        self._node2_obs: ParamAnomalyObservation | None = None
        self._node2_gt: GroundTruthParamAnnotation | None = None

        self._subenv3_case: dict[str, Any] | None = None
        self._subenv3_gt: GroundTruthBehavioralAnnotation | None = None

    def _load_cases(self, filename: str) -> list[dict[str, Any]]:
        path = _TEST_SET_DIR / filename
        if not path.exists():
            log.warning("Missing test case file: %s", path)
            return []
        with path.open("r", encoding="utf-8") as f:
            payload = json.load(f)
        cases = payload.get("cases", [])
        return cases if isinstance(cases, list) else []

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> Observation:
        if not self._subenv1_cases:
            raise RuntimeError("subenv1_cases.json has no cases; cannot reset environment")

        self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
        self._scores = {}
        self._phase = "node1"

        self._node1_action = None
        self._node2_obs = None
        self._node2_gt = None

        if seed is not None:
            rng = random.Random(seed)
            idx = rng.randrange(len(self._subenv1_cases))
        else:
            idx = self._cursor % len(self._subenv1_cases)
            self._cursor += 1

        self._current_case = self._subenv1_cases[idx]
        self._node1_obs = ImageDiagnosticsObservation.model_validate(
            self._current_case["observation"]
        )
        self._node1_gt = GroundTruthImageAnnotation.model_validate(
            self._current_case["ground_truth"]
        )

        self._subenv3_case = self._pick_subenv3_case(idx)
        self._subenv3_gt = self._resolve_behavioral_ground_truth(self._subenv3_case)

        return Observation(
            done=False,
            reward=0.0,
            metadata={
                "node": "node1_image_diagnostician",
                "step": 0,
                "episode_id": self._state.episode_id,
                "case_id": self._current_case.get("id", "unknown"),
                "instruction": (
                    "Diagnose reference-image quality and prompt risks from the provided "
                    "signals."
                ),
                "observation": self._node1_obs.model_dump(mode="json"),
                "expected_action_schema": "ImageDiagnosticsAction",
            },
        )

    def step(
        self,
        action: Action,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> Observation:
        self._state.step_count += 1

        if self._phase == "node1":
            return self._handle_node1_action(action)
        if self._phase == "node2":
            return self._handle_node2_action(action)
        if self._phase == "node8":
            return self._handle_node8_action(action)

        return Observation(
            done=True,
            reward=0.0,
            metadata={"error": "Episode complete. Call reset() for a new episode."},
        )

    def _handle_node1_action(self, action: Action) -> Observation:
        if self._current_case is None or self._node1_obs is None or self._node1_gt is None:
            raise RuntimeError("Environment is not initialized. Call reset() before step().")

        node1_action = self._parse_action(action, ImageDiagnosticsAction)
        self._node1_action = node1_action

        node1_score = grade_image_diagnostics(node1_action, self._node1_gt)
        self._scores["node1"] = node1_score

        self._node2_obs = self._build_param_anomaly_observation(node1_action, self._node1_obs)
        self._node2_gt = self._resolve_param_ground_truth(self._current_case, self._node2_obs)
        self._phase = "node2"

        return Observation(
            done=False,
            reward=node1_score,
            metadata={
                "node": "node2_param_anomaly_detector",
                "step": 1,
                "node1_score": node1_score,
                "instruction": (
                    "Analyze proposed generation parameters, detect anomalies, and propose "
                    "directional fixes."
                ),
                "observation": self._node2_obs.model_dump(mode="json"),
                "expected_action_schema": "ParamAnomalyAction",
            },
        )

    def _handle_node2_action(self, action: Action) -> Observation:
        if (
            self._node1_action is None
            or self._node1_gt is None
            or self._node2_obs is None
            or self._node2_gt is None
        ):
            raise RuntimeError("Node 2 prerequisites missing. Run Node 1 step first.")

        node2_action = self._parse_action(action, ParamAnomalyAction)
        node2_score = grade_anomaly_detection(node2_action, self._node2_gt)
        self._scores["node2"] = node2_score

        subenv1_score = 0.50 * self._scores.get("node1", 0.0) + 0.50 * node2_score
        self._scores["subenv1"] = subenv1_score

        reference_handoff = produce_reference_audit_handoff(
            self._node1_action,
            node2_action,
            self._node1_gt,
            self._node2_gt,
        )
        subenv2_score = self._run_subenv2_internally(
            reference_handoff.risk_profile,
            reference_handoff.estimated_drift_risk,
        )
        self._scores["subenv2"] = subenv2_score

        phoneme_obs = self._build_phoneme_risk_observation()
        self._phase = "node8"

        return Observation(
            done=False,
            reward=subenv1_score,
            metadata={
                "node": "node8_phoneme_risk_assessor",
                "step": 2,
                "subenv1_score": subenv1_score,
                "subenv2_score": subenv2_score,
                "instruction": (
                    "Assess phoneme-level behavioral risk from weight evidence and provide "
                    "mitigation guidance."
                ),
                "observation": phoneme_obs.model_dump(mode="json"),
                "expected_action_schema": "PhonemeRiskAction",
            },
        )

    def _handle_node8_action(self, action: Action) -> Observation:
        if self._subenv3_gt is None:
            raise RuntimeError("Behavioral ground truth is unavailable. Reset the environment.")

        node8_action = self._parse_action(action, PhonemeRiskAction)
        node8_score = grade_behavioral_audit(node8_action, self._subenv3_gt)

        self._scores["node8"] = node8_score
        self._scores["subenv3"] = node8_score

        s1 = self._scores.get("subenv1", 0.0)
        s2 = self._scores.get("subenv2", 0.0)
        s3 = self._scores.get("subenv3", 0.0)
        final_score = 0.25 * s1 + 0.35 * s2 + 0.40 * s3
        self._scores["final"] = final_score

        self._phase = "done"

        return Observation(
            done=True,
            reward=final_score,
            metadata={
                "node": "episode_complete",
                "step": 3,
                "reward_formula": "0.25 * subenv1 + 0.35 * subenv2 + 0.40 * subenv3",
                "scores": {
                    "subenv1_score": s1,
                    "subenv2_score": s2,
                    "subenv3_score": s3,
                    "final_score": final_score,
                    "breakdown": self._scores,
                },
            },
        )

    def _parse_action(self, action: Action, expected_type: type[_T]) -> _T:
        if isinstance(action, expected_type):
            return action

        payload: Any = None

        if isinstance(action, dict):
            payload = action.get("action", action)
        elif isinstance(action, Action):
            payload = action.metadata.get("action", action.metadata)
        elif hasattr(action, "metadata") and isinstance(action.metadata, dict):
            payload = action.metadata.get("action", action.metadata)
        elif hasattr(action, "model_dump"):
            payload = action.model_dump()

        if isinstance(payload, str):
            payload = json.loads(payload)

        if not isinstance(payload, dict):
            raise ValueError(
                f"Unable to parse action payload for {expected_type.__name__}; "
                f"received {type(action).__name__}"
            )

        return expected_type.model_validate(payload)

    def _pick_subenv3_case(self, idx: int) -> dict[str, Any] | None:
        if not self._subenv3_cases:
            return None
        return self._subenv3_cases[idx % len(self._subenv3_cases)]

    def _resolve_behavioral_ground_truth(
        self,
        case: dict[str, Any] | None,
    ) -> GroundTruthBehavioralAnnotation:
        if case and isinstance(case.get("ground_truth"), dict):
            try:
                return GroundTruthBehavioralAnnotation.model_validate(case["ground_truth"])
            except Exception as exc:  # noqa: BLE001
                log.warning("Failed to parse behavioral ground truth from case: %s", exc)

        return GroundTruthBehavioralAnnotation(
            phoneme_risk_ranking=[],
            predicted_behavior_triggers=[],
            risky_phoneme_clusters=[],
            model_behavioral_safety="safe",
            valid_mitigation_set=set(),
        )

    def _build_param_anomaly_observation(
        self,
        node1_action: ImageDiagnosticsAction,
        image_obs: ImageDiagnosticsObservation,
    ) -> ParamAnomalyObservation:
        return ParamAnomalyObservation(
            proposed_config={"cfg": 7.5, "denoise_alt": 0.30, "eta": 0.10},
            regime=node1_action.regime_classification,
            identified_risk_factors=node1_action.identified_risk_factors,
            image_usability_score=node1_action.image_usability_score,
            face_occupancy_ratio=image_obs.face_occupancy_ratio,
            estimated_yaw_degrees=image_obs.estimated_yaw_degrees,
            background_complexity_score=image_obs.background_complexity_score,
            lighting_uniformity_score=image_obs.lighting_uniformity_score,
            occlusion_detected=image_obs.occlusion_detected,
            prompt_identity_anchoring=image_obs.identity_anchoring_strength,
            prompt_token_count=image_obs.prompt_token_count,
            conflicting_descriptors=image_obs.conflicting_descriptors,
        )

    def _resolve_param_ground_truth(
        self,
        case: dict[str, Any],
        node2_obs: ParamAnomalyObservation,
    ) -> GroundTruthParamAnnotation:
        for key in ("param_ground_truth", "ground_truth_param", "node2_ground_truth"):
            if isinstance(case.get(key), dict):
                return GroundTruthParamAnnotation.model_validate(case[key])

        # Current test sets do not include explicit Node 2 ground truth. For
        # deterministic grading, we synthesize a reference annotation from the
        # built-in Node 2 heuristic on the same observation.
        baseline = detect_param_anomalies(node2_obs)
        return GroundTruthParamAnnotation(
            config_risk_level=baseline.config_risk_level,
            anomalies=baseline.anomalies,
            predicted_failure_modes=list(baseline.predicted_failure_modes),
            valid_fix_directions=baseline.directional_fixes,
        )

    def _build_clip_evidence_dossier(self, obs: ClipSignalObservation) -> ClipEvidenceDossier:
        drift = obs.identity_cosine_drift
        if drift < 0.05:
            drift_severity = "none"
        elif drift < 0.12:
            drift_severity = "minor"
        elif drift < 0.22:
            drift_severity = "moderate"
        else:
            drift_severity = "severe"

        temporal_instability = (
            obs.landmark_stability_score > 0.04 or obs.frame_difference_mean > 15.0
        )

        lip_sync = obs.lip_sync_confidence
        if lip_sync >= 0.75:
            lip_sync_quality = "good"
        elif lip_sync >= 0.50:
            lip_sync_quality = "acceptable"
        elif lip_sync >= 0.20:
            lip_sync_quality = "poor"
        else:
            lip_sync_quality = "absent"

        dataset_redundancy = min(1.0, obs.similar_clips_accepted / 10.0)

        if drift_severity in ("none", "minor") and not temporal_instability and lip_sync >= 0.50:
            training_impact = "positive"
        elif drift_severity == "severe" or (temporal_instability and lip_sync < 0.20):
            training_impact = "negative"
        else:
            training_impact = "neutral"

        primary_rejection_reason = None
        if drift_severity in ("moderate", "severe"):
            primary_rejection_reason = (
                f"identity cosine drift {drift:.3f} exceeds threshold; severity={drift_severity}"
            )
        elif temporal_instability and lip_sync < 0.20:
            primary_rejection_reason = "temporal instability with absent lip sync"

        summary = (
            f"Drift={drift_severity}, temporal_instability={temporal_instability}, "
            f"lip_sync={lip_sync_quality}, phoneme_value={obs.phoneme_coverage_new:.2f}, "
            f"redundancy={dataset_redundancy:.2f}."
        )

        return ClipEvidenceDossier(
            clip_id=obs.clip_id,
            identity_drift_severity=drift_severity,
            temporal_instability_flag=temporal_instability,
            lip_sync_quality=lip_sync_quality,
            unique_phoneme_value=obs.phoneme_coverage_new,
            dataset_redundancy_score=dataset_redundancy,
            estimated_training_impact=training_impact,
            primary_rejection_reason=primary_rejection_reason,
            evidence_summary=summary,
        )

    def _run_subenv2_internally(self, reference_risk_profile: str, estimated_drift_risk: float) -> float:
        if not self._subenv2_cases:
            return 0.5

        scores: list[float] = []

        for clip_case in self._subenv2_cases[:5]:
            try:
                clip_obs = ClipSignalObservation.model_validate(clip_case["observation"])
                clip_gt = GroundTruthClipAnnotation.model_validate(clip_case["ground_truth"])

                dossier = self._build_clip_evidence_dossier(clip_obs)
                damage_map = {"none": 0.0, "minor": 0.05, "moderate": 0.20, "severe": 0.50}
                marginal_damage = damage_map.get(dossier.identity_drift_severity, 0.1)

                clip_disp_obs = ClipDispositionObservation(
                    evidence_dossier=dossier,
                    minimum_clips_needed=20,
                    phoneme_gap_severity={},
                    pose_gap_severity={},
                    budget_remaining=max(0, 50 - clip_obs.clips_audited_so_far),
                    reference_risk_profile=reference_risk_profile,
                    estimated_drift_risk=estimated_drift_risk,
                    marginal_training_damage=marginal_damage,
                    marginal_coverage_gain=float(clip_obs.phoneme_coverage_new),
                )

                clip_action = recommend_clip_disposition(clip_disp_obs)
                scores.append(grade_clip_disposition(clip_action, clip_gt))
            except Exception as exc:  # noqa: BLE001
                log.warning("Sub-env 2 internal clip evaluation failed: %s", exc)

        return sum(scores) / len(scores) if scores else 0.5

    def _build_phoneme_risk_observation(self) -> PhonemeRiskObservation:
        if self._subenv3_case and isinstance(self._subenv3_case.get("observation"), dict):
            return PhonemeRiskObservation.model_validate(self._subenv3_case["observation"])

        if self._subenv3_case and isinstance(self._subenv3_case.get("phoneme_obs"), dict):
            return PhonemeRiskObservation.model_validate(self._subenv3_case["phoneme_obs"])

        return PhonemeRiskObservation.model_validate(
            {
                "weight_evidence": {
                    "weight_file_id": "fallback.safetensors",
                    "training_quality": "healthy",
                    "rank_utilization_assessment": "efficient",
                    "high_entropy_token_flags": [],
                    "layer_anomaly_flags": [],
                    "overall_behavioral_risk": "low",
                    "evidence_summary": "No significant anomalies detected.",
                },
                "high_entropy_token_flags": [],
                "phoneme_vocabulary": [
                    "AA",
                    "AE",
                    "AH",
                    "AO",
                    "AW",
                    "AY",
                    "EH",
                    "ER",
                    "EY",
                    "IH",
                    "IY",
                    "OW",
                    "OY",
                    "UH",
                    "UW",
                ],
                "phoneme_to_token_indices": {},
                "phoneme_entropy_scores": {},
                "phoneme_influence_scores": {},
                "phoneme_cooccurrence_anomalies": [],
                "behavior_vocabulary": ["smile", "blink", "head_turn"],
                "training_data_phoneme_distribution": None,
                "suspected_anomalous_phonemes_from_subenv2": None,
            }
        )

    @property
    def state(self) -> State:
        return self._state