env/env.py

"""Core RecallTrace environment with deterministic action execution."""

from __future__ import annotations

from copy import deepcopy
from typing import Any, Dict, Tuple

from env.models import EnvironmentState, InspectionEvidence, RecallAction, RecallObservation, RewardSignal, StepInfo, TaskDefinition
from scenario.scenario import build_scenario, list_task_specs


class RecallTraceEnv:
    """Deterministic OpenEnv-style environment for product recall containment."""

    ACTIONS = [
        "inspect_node",
        "trace_lot",
        "quarantine",
        "notify",
        "finalize",
    ]

    def __init__(
        self,
        scenario_data: Dict[str, Any] | None = None,
        task_id: str | None = None,
        phase: int | None = 1,
    ):
        self._scenario_template = deepcopy(scenario_data) if scenario_data is not None else build_scenario(task_id=task_id, phase=phase)
        self.task = self._build_task_definition(self._scenario_template)
        self.state_data: Dict[str, Any] = {}
        self.ground_truth: Dict[str, Any] = {}
        self.done = False
        self.last_reward = RewardSignal(value=0.0, reason="Environment initialized.", components={})

    @classmethod
    def available_tasks(cls) -> list[TaskDefinition]:
        return [TaskDefinition(**task_spec) for task_spec in list_task_specs()]

    def reset(self, task_id: str | None = None, phase: int | None = None) -> RecallObservation:
        """Start a new deterministic scenario and recompute ground truth."""
        if task_id is not None or phase is not None:
            self._scenario_template = build_scenario(task_id=task_id, phase=phase)
            self.task = self._build_task_definition(self._scenario_template)

        self.done = False
        self.last_reward = RewardSignal(value=0.0, reason="Episode reset.", components={})

        scenario = deepcopy(self._scenario_template)
        self.state_data = {
            "task_id": scenario["task_id"],
            "phase": scenario["phase"],
            "recall_notice": scenario["recall_notice"],
            "contaminated_lot_hint": scenario["contaminated_lot"],
            "shipment_graph": scenario["shipment_graph"],
            "lot_catalog": scenario["lot_catalog"],
            "nodes": scenario["nodes"],
            "history": [],
            "discovered_shipments": {},
            "inspected_nodes": set(),
            "inspection_results": {},
            "traced_lots": {},
            "notified_nodes": set(),
            "quarantine_log": [],
            "steps_taken": 0,
            "max_steps": scenario["max_steps"],
        }
        self.ground_truth = self._build_ground_truth(scenario)
        return self._get_observation()

    def step(self, action: RecallAction | Dict[str, Any]) -> Tuple[RecallObservation, float, bool, Dict[str, Any]]:
        """Execute an action and return observation, reward, done, info."""
        if self.done:
            return self._get_observation(), 0.0, True, {
                "message": "Environment already finalized.",
                "action_type": "noop",
                "reward_breakdown": {},
            }

        validated_action = action if isinstance(action, RecallAction) else RecallAction.model_validate(action)
        self.state_data["steps_taken"] += 1

        handler = getattr(self, f"_handle_{validated_action.type.value}")
        reward_signal, info = handler(validated_action)
        self.last_reward = reward_signal

        if not self.done and self.state_data["steps_taken"] >= self.state_data["max_steps"]:
            self.done = True
            timeout_penalty = -0.25
            reward_signal = RewardSignal(
                value=max(-1.0, reward_signal.value + timeout_penalty),
                reason="Step budget exhausted before finalizing containment.",
                components={**reward_signal.components, "timeout_penalty": timeout_penalty},
            )
            info = {
                **info,
                "message": "Step budget exhausted before finalizing containment.",
                "reward_breakdown": reward_signal.components,
            }
            self._record_history("Episode terminated after exhausting the step budget")
            self.last_reward = reward_signal

        return self._get_observation(), reward_signal.value, self.done, info

    def state(self) -> EnvironmentState:
        """Return the full internal state for debugging and graders."""
        return EnvironmentState(
            done=self.done,
            task=self.task,
            steps_taken=self.state_data.get("steps_taken", 0),
            state_data=deepcopy(self._serialize_state(self.state_data)),
            ground_truth=deepcopy(self.ground_truth),
        )

    def _get_observation(self) -> RecallObservation:
        return RecallObservation(
            task_id=self.state_data["task_id"],
            phase=self.state_data["phase"],
            recall_notice=self.state_data["recall_notice"],
            available_actions=list(self.ACTIONS),
            inventory=self._inventory_snapshot(),
            discovered_shipments=deepcopy(self.state_data["discovered_shipments"]),
            inspected_nodes=sorted(self.state_data["inspected_nodes"]),
            inspection_results=deepcopy(self.state_data["inspection_results"]),
            trace_results=deepcopy(self.state_data["traced_lots"]),
            notified_nodes=sorted(self.state_data["notified_nodes"]),
            quarantined_inventory=self._quarantine_snapshot(),
            history=list(self.state_data["history"]),
            steps_taken=self.state_data["steps_taken"],
            remaining_step_budget=max(0, self.state_data["max_steps"] - self.state_data["steps_taken"]),
        )

    def _handle_inspect_node(self, action: RecallAction) -> tuple[RewardSignal, Dict[str, Any]]:
        node_id = self._require_node(action.node_id)
        node = self.state_data["nodes"][node_id]
        repeated = node_id in self.state_data["inspected_nodes"]

        self.state_data["inspected_nodes"].add(node_id)
        self.state_data["discovered_shipments"][node_id] = list(self.state_data["shipment_graph"].get(node_id, []))
        findings = {
            lot_id: InspectionEvidence.model_validate(payload)
            for lot_id, payload in node.get("inspection_findings", {}).items()
        }
        self.state_data["inspection_results"][node_id] = findings
        self._record_history(f"Inspected node {node_id}")

        unsafe_total = sum(item.unsafe_quantity for item in findings.values())
        value = -0.03 if repeated else 0.08 + min(0.12, unsafe_total / 500.0)
        reason = "Repeated inspection provided no new information." if repeated else "Inspection revealed inventory evidence."
        reward = RewardSignal(
            value=round(value, 4),
            reason=reason,
            components={
                "inspection_value": round(value, 4),
            },
        )
        info = StepInfo(
            message=f"Inspected node {node_id} and collected node evidence.",
            action_type=action.type.value,
            reward_breakdown=reward.components,
        ).model_dump()
        info.update(
            {
                "node_id": node_id,
                "inventory": deepcopy(node["inventory"]),
                "quarantined_inventory": deepcopy(node["quarantined_inventory"]),
                "outbound_shipments": list(self.state_data["shipment_graph"].get(node_id, [])),
                "inspection_findings": {lot_id: item.model_dump() for lot_id, item in findings.items()},
            }
        )
        return reward, info

    def _handle_trace_lot(self, action: RecallAction) -> tuple[RewardSignal, Dict[str, Any]]:
        lot_id = action.lot_id
        if not lot_id:
            raise ValueError("trace_lot action requires 'lot_id'.")

        traced_lots = self._resolve_related_lots(lot_id)
        impacted_nodes = []
        impacted_quantities = {}
        impacted_lots = {}
        discovered_nodes = 0

        for node_id, node_data in self.state_data["nodes"].items():
            node_total = 0
            node_lots = []
            for candidate_lot in traced_lots:
                available_qty = node_data["inventory"].get(candidate_lot, 0)
                quarantined_qty = node_data["quarantined_inventory"].get(candidate_lot, 0)
                total_qty = available_qty + quarantined_qty
                if total_qty > 0:
                    node_total += total_qty
                    node_lots.append(candidate_lot)
            if node_total > 0:
                impacted_nodes.append(node_id)
                impacted_quantities[node_id] = node_total
                impacted_lots[node_id] = node_lots
                if node_id not in self.state_data["discovered_shipments"]:
                    discovered_nodes += 1

        self.state_data["traced_lots"][lot_id] = {
            "root_lot": self._root_lot_for(lot_id),
            "matched_lots": sorted(traced_lots),
            "affected_nodes": impacted_nodes,
            "lots_by_node": impacted_lots,
            "quantities_by_node": impacted_quantities,
        }
        self._record_history(f"Traced lot {lot_id} across {', '.join(sorted(traced_lots))}")

        if not impacted_nodes:
            reward_value = -0.1
            reason = "Trace returned no impacted nodes."
        elif self._root_lot_for(lot_id) in self.ground_truth["affected_roots"]:
            reward_value = 0.12 + min(0.13, discovered_nodes * 0.03 + len(traced_lots) * 0.02)
            reason = "Trace identified the affected lineage across the network."
        else:
            reward_value = 0.02
            reason = "Trace ran, but the lot is outside the affected lineage."

        reward = RewardSignal(
            value=round(reward_value, 4),
            reason=reason,
            components={
                "trace_value": round(reward_value, 4),
            },
        )
        info = StepInfo(
            message=f"Traced lot {lot_id} across the shipment network.",
            action_type=action.type.value,
            reward_breakdown=reward.components,
        ).model_dump()
        info.update(
            {
                "lot_id": lot_id,
                "root_lot": self._root_lot_for(lot_id),
                "matched_lots": sorted(traced_lots),
                "affected_nodes": impacted_nodes,
                "lots_by_node": impacted_lots,
                "quantities_by_node": impacted_quantities,
                "total_quantity": sum(impacted_quantities.values()),
            }
        )
        return reward, info

    def _handle_quarantine(self, action: RecallAction) -> tuple[RewardSignal, Dict[str, Any]]:
        node_id = self._require_node(action.node_id)
        lot_id = action.lot_id
        if not lot_id:
            raise ValueError("quarantine action requires 'lot_id'.")

        node = self.state_data["nodes"][node_id]
        available_qty = node["inventory"].get(lot_id, 0)
        if available_qty <= 0:
            reward = RewardSignal(
                value=-0.2,
                reason="Attempted to quarantine stock that is not available.",
                components={"invalid_quarantine": -0.2},
            )
            self._record_history(f"Failed quarantine for {lot_id} at {node_id}: no available stock")
            info = StepInfo(
                message="No available stock to quarantine.",
                action_type=action.type.value,
                reward_breakdown=reward.components,
            ).model_dump()
            info.update({"node_id": node_id, "lot_id": lot_id})
            return reward, info

        requested_qty = action.quantity or available_qty
        quarantined_qty = min(requested_qty, available_qty)
        node["inventory"][lot_id] = available_qty - quarantined_qty
        if node["inventory"][lot_id] == 0:
            del node["inventory"][lot_id]
        node["quarantined_inventory"][lot_id] = node["quarantined_inventory"].get(lot_id, 0) + quarantined_qty

        self.state_data["quarantine_log"].append({"node_id": node_id, "lot_id": lot_id, "quantity": quarantined_qty})
        self._record_history(f"Quarantined {quarantined_qty} units of {lot_id} at {node_id}")

        correct_qty = self.ground_truth["correct_quantities"].get(node_id, {}).get(lot_id, 0)
        cumulative_quarantined = node["quarantined_inventory"].get(lot_id, 0)
        delta = cumulative_quarantined - correct_qty

        if correct_qty == 0:
            reward_value = -0.35
            reason = "Quarantined safe inventory outside the recall scope."
        elif delta == 0:
            reward_value = 0.28
            reason = "Quarantine exactly matched the unsafe quantity."
        elif delta < 0:
            reward_value = max(0.05, 0.22 * (cumulative_quarantined / correct_qty))
            reason = "Quarantine made partial progress but missed some unsafe stock."
        else:
            reward_value = max(-0.25, -0.08 * delta)
            reason = "Quarantine overreached and blocked safe inventory."

        reward = RewardSignal(
            value=round(reward_value, 4),
            reason=reason,
            components={
                "quarantine_value": round(reward_value, 4),
                "target_quantity": float(correct_qty),
                "quarantined_quantity": float(cumulative_quarantined),
            },
        )
        info = StepInfo(
            message=f"Updated quarantine for {lot_id} at {node_id}.",
            action_type=action.type.value,
            reward_breakdown=reward.components,
        ).model_dump()
        info.update(
            {
                "node_id": node_id,
                "lot_id": lot_id,
                "quarantined_quantity": quarantined_qty,
                "remaining_inventory": node["inventory"].get(lot_id, 0),
                "cumulative_quarantined": cumulative_quarantined,
                "target_contaminated_quantity": correct_qty,
            }
        )
        return reward, info

    def _handle_notify(self, action: RecallAction) -> tuple[RewardSignal, Dict[str, Any]]:
        requested_target = action.node_id or "all"
        if requested_target in ("all", "all_nodes"):
            targets = list(self.state_data["nodes"].keys())
        else:
            targets = [self._require_node(requested_target)]

        newly_notified = []
        for node_id in targets:
            if node_id not in self.state_data["notified_nodes"]:
                self.state_data["notified_nodes"].add(node_id)
                newly_notified.append(node_id)

        affected_newly_notified = sum(1 for node_id in newly_notified if node_id in self.ground_truth["affected_nodes"])
        unaffected_newly_notified = len(newly_notified) - affected_newly_notified

        if not newly_notified:
            reward_value = -0.05
            reason = "Notification repeated without adding new recipients."
        else:
            reward_value = min(0.18, affected_newly_notified * 0.04) - unaffected_newly_notified * 0.01
            reason = "Notifications dispatched to downstream stakeholders."

        reward = RewardSignal(
            value=round(reward_value, 4),
            reason=reason,
            components={
                "notification_value": round(reward_value, 4),
            },
        )
        if newly_notified:
            self._record_history(f"Sent notifications to {', '.join(newly_notified)}")
        else:
            self._record_history("Notification action repeated without new recipients")

        info = StepInfo(
            message="Processed notification action.",
            action_type=action.type.value,
            reward_breakdown=reward.components,
        ).model_dump()
        info.update({"notified_nodes": targets, "newly_notified": newly_notified})
        return reward, info

    def _handle_finalize(self, action: RecallAction) -> tuple[RewardSignal, Dict[str, Any]]:
        del action
        self.done = True
        quarantine_match = self._compute_quarantine_match()

        missing_quantity_total = sum(
            quantity
            for lot_quantities in quarantine_match["missing_quantities"].values()
            for quantity in lot_quantities.values()
        )
        over_quantity_total = sum(
            quantity
            for lot_quantities in quarantine_match["over_quarantined_quantities"].values()
            for quantity in lot_quantities.values()
        )
        total_affected_quantity = self.ground_truth["total_affected_quantity"] or 1
        quarantine_score = max(0.0, 1.0 - ((missing_quantity_total + (1.25 * over_quantity_total)) / total_affected_quantity))

        notified_affected_nodes = set(self.ground_truth["affected_nodes"]).intersection(self.state_data["notified_nodes"])
        affected_node_total = len(self.ground_truth["affected_nodes"]) or 1
        notification_score = len(notified_affected_nodes) / affected_node_total

        investigated_nodes = set(self.state_data["inspected_nodes"]).intersection(self.ground_truth["affected_nodes"])
        investigation_score = len(investigated_nodes) / affected_node_total

        efficiency_penalty_steps = max(0, self.state_data["steps_taken"] - max(4, affected_node_total + 3))
        efficiency_score = max(0.0, 1.0 - (efficiency_penalty_steps / self.state_data["max_steps"]))

        score = round(
            (0.55 * quarantine_score) + (0.2 * notification_score) + (0.15 * investigation_score) + (0.1 * efficiency_score),
            4,
        )

        reward = RewardSignal(
            value=score,
            reason="Final recall response scored.",
            components={
                "quarantine_score": round(quarantine_score, 4),
                "notification_score": round(notification_score, 4),
                "investigation_score": round(investigation_score, 4),
                "efficiency_score": round(efficiency_score, 4),
            },
        )
        self._record_history("Finalized recall response")

        info = StepInfo(
            message="Finalized recall response.",
            action_type="finalize",
            score=score,
            reward_breakdown=reward.components,
        ).model_dump()
        info.update(
            {
                "score": score,
                "quarantine_score": round(quarantine_score, 4),
                "notification_score": round(notification_score, 4),
                "investigation_score": round(investigation_score, 4),
                "efficiency_score": round(efficiency_score, 4),
                "all_affected_nodes_notified": notification_score == 1.0,
                "all_affected_stock_quarantined": missing_quantity_total == 0 and over_quantity_total == 0,
                "quarantine_match": quarantine_match,
            }
        )
        return reward, info

    def _build_ground_truth(self, scenario: Dict[str, Any]) -> Dict[str, Any]:
        contaminated_roots = {
            self._root_lot_for(lot_id, scenario["lot_catalog"])
            for lot_id, lot_data in scenario["lot_catalog"].items()
            if lot_data.get("contaminated", False)
        }

        correct_quantities: Dict[str, Dict[str, int]] = {}
        affected_nodes = set()
        affected_lots = set()

        for node_id, node_data in scenario["nodes"].items():
            for lot_id, finding in node_data.get("inspection_findings", {}).items():
                unsafe_quantity = int(finding.get("unsafe_quantity", 0))
                if unsafe_quantity > 0:
                    affected_nodes.add(node_id)
                    affected_lots.add(lot_id)
                    correct_quantities.setdefault(node_id, {})[lot_id] = unsafe_quantity

        total_affected_quantity = sum(
            quantity
            for node_quantities in correct_quantities.values()
            for quantity in node_quantities.values()
        )
        return {
            "affected_lots": sorted(affected_lots),
            "affected_nodes": sorted(affected_nodes),
            "affected_roots": sorted(contaminated_roots),
            "correct_quantities": correct_quantities,
            "total_affected_quantity": total_affected_quantity,
        }

    def _compute_quarantine_match(self) -> Dict[str, Any]:
        missing_quantities: Dict[str, Dict[str, int]] = {}
        over_quarantined_quantities: Dict[str, Dict[str, int]] = {}

        for node_id, node_data in self.state_data["nodes"].items():
            expected = self.ground_truth["correct_quantities"].get(node_id, {})
            actual = node_data["quarantined_inventory"]
            relevant_lots = set(expected) | set(actual)

            for lot_id in relevant_lots:
                expected_qty = expected.get(lot_id, 0)
                actual_qty = actual.get(lot_id, 0)
                if actual_qty < expected_qty:
                    missing_quantities.setdefault(node_id, {})[lot_id] = expected_qty - actual_qty
                elif actual_qty > expected_qty:
                    over_quarantined_quantities.setdefault(node_id, {})[lot_id] = actual_qty - expected_qty

        return {
            "missing_quantities": missing_quantities,
            "over_quarantined_quantities": over_quarantined_quantities,
        }

    def _inventory_snapshot(self) -> Dict[str, Dict[str, int]]:
        return {node_id: deepcopy(node_data["inventory"]) for node_id, node_data in self.state_data["nodes"].items()}

    def _quarantine_snapshot(self) -> Dict[str, Dict[str, int]]:
        return {
            node_id: deepcopy(node_data["quarantined_inventory"])
            for node_id, node_data in self.state_data["nodes"].items()
            if node_data["quarantined_inventory"]
        }

    def _resolve_related_lots(self, lot_id: str) -> set[str]:
        root_lot = self._root_lot_for(lot_id)
        return {
            candidate_lot
            for candidate_lot in self.state_data["lot_catalog"].keys()
            if self._root_lot_for(candidate_lot) == root_lot or candidate_lot == lot_id
        }

    def _root_lot_for(self, lot_id: str, lot_catalog: Dict[str, Dict[str, Any]] | None = None) -> str:
        catalog = lot_catalog or self.state_data.get("lot_catalog", {})
        if lot_id not in catalog:
            return lot_id
        return catalog[lot_id].get("root_lot", lot_id)

    def _build_task_definition(self, scenario: Dict[str, Any]) -> TaskDefinition:
        return TaskDefinition(
            task_id=scenario["task_id"],
            name=scenario["name"],
            difficulty=scenario["difficulty"],
            objective=scenario["objective"],
            max_steps=scenario["max_steps"],
        )

    def _require_node(self, node_id: str | None) -> str:
        if not node_id:
            raise ValueError("Action requires 'node_id'.")
        if node_id not in self.state_data["nodes"]:
            raise ValueError(f"Unknown node_id '{node_id}'.")
        return node_id

    def _record_history(self, message: str) -> None:
        self.state_data["history"].append(message)

    def _serialize_state(self, value: Any) -> Any:
        if isinstance(value, dict):
            return {key: self._serialize_state(item) for key, item in value.items()}
        if isinstance(value, set):
            return sorted(value)
        if isinstance(value, list):
            return [self._serialize_state(item) for item in value]
        if hasattr(value, "model_dump"):
            return value.model_dump()
        return value