"""Environment-level roundtrip serialization tests for Search & Rescue.

Tests verify that full S&R environment state can be saved and restored via
CoGridEnv.get_state/set_state, maintaining byte-perfect fidelity.

This parallels the Overcooked environment tests in test_state_serialization.py.

Requirements verified:
- Environment roundtrip preserves timestep, agent positions, grid state
- RedVictim mid-rescue state survives full environment roundtrip
- Agent inventory is preserved across roundtrip
- RNG state is preserved for reproducibility
- Termination state is preserved after restoration
"""

import pytest
import numpy as np

from cogrid.core.actions import Actions
from cogrid.envs import registry
from cogrid.envs.search_rescue.search_rescue_grid_objects import RedVictim, MedKit


class TestSearchRescueEnvSerialization:
    """Test state serialization for Search & Rescue environments."""

    @pytest.fixture
    def env(self):
        """Create a Search & Rescue test environment."""
        env = registry.make("SearchRescue-Test-V0")
        return env

    def test_basic_get_state(self, env):
        """Verify get_state returns dict with expected keys."""
        env.reset(seed=42)

        state = env.get_state()

        # Verify top-level keys present
        assert "timestep" in state
        assert "grid" in state
        assert "agents" in state
        assert "rng_state" in state
        assert "config" in state

        # Verify grid subkeys
        assert "encoding" in state["grid"]
        assert "object_metadata" in state["grid"]

        # Verify agents are dict keyed by agent_id
        assert isinstance(state["agents"], dict)
        assert len(state["agents"]) == 2  # SearchRescue-Test-V0 has 2 agents

    def test_roundtrip_after_reset(self, env):
        """Save/restore immediately after reset, verify agent positions and grid match."""
        env.reset(seed=42)

        # Record original state
        original_positions = {
            agent_id: tuple(agent.pos) for agent_id, agent in env.env_agents.items()
        }
        original_grid_encoding = env.grid.encode(scope=env.scope).copy()

        # Save state
        state = env.get_state()

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Verify timestep
        assert env2.t == 0

        # Verify agent positions
        for agent_id, agent in env2.env_agents.items():
            assert tuple(agent.pos) == original_positions[agent_id]

        # Verify grid encoding matches exactly
        np.testing.assert_array_equal(
            original_grid_encoding, env2.grid.encode(scope=env2.scope)
        )

    def test_roundtrip_after_steps(self, env):
        """Run 10 steps, save, restore, verify timestep and grid encoding match."""
        env.reset(seed=42)

        # Take 10 steps with actions
        actions = {0: Actions.MoveRight, 1: Actions.MoveLeft}
        for _ in range(10):
            env.step(actions)

        # Save state after 10 steps
        state = env.get_state()

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Verify timestep preserved
        assert env.t == env2.t == 10

        # Verify grid encoding matches exactly
        np.testing.assert_array_equal(
            env.grid.encode(scope=env.scope), env2.grid.encode(scope=env2.scope)
        )

        # Verify agent positions match
        for agent_id in env.env_agents:
            np.testing.assert_array_equal(
                env.env_agents[agent_id].pos, env2.env_agents[agent_id].pos
            )

    def test_redvictim_mid_rescue_roundtrip(self, env):
        """Set up RedVictim with active countdown, save state mid-rescue, restore to new env.

        The two-step rescue mechanic:
        1. Medic/agent with MedKit toggles RedVictim, starting 30-step countdown
        2. Different agent must toggle within countdown window to complete rescue

        This test verifies that countdown and first_toggle_agent survive roundtrip.
        """
        env.reset(seed=42)

        # Find a RedVictim in the grid (layout has 'R' character if present)
        # Since the test layout may not have RedVictim, we'll manually add one
        # Find an empty floor space
        red_victim_pos = None
        for row in range(env.grid.height):
            for col in range(env.grid.width):
                obj = env.grid.get(row, col)
                if obj is None:  # Empty floor
                    red_victim_pos = (row, col)
                    break
            if red_victim_pos:
                break

        # Create and place a RedVictim with active rescue state
        red_victim = RedVictim(state=0)
        red_victim.toggle_countdown = 15
        red_victim.first_toggle_agent = "agent_0"
        env.grid.set(*red_victim_pos, red_victim)

        # Save state
        state = env.get_state()

        # Verify the state contains the RedVictim extra state
        pos_key = f"{red_victim_pos[0]},{red_victim_pos[1]}"
        assert pos_key in state["grid"]["object_metadata"]
        metadata = state["grid"]["object_metadata"][pos_key]
        assert metadata["extra_state"]["toggle_countdown"] == 15
        assert metadata["extra_state"]["first_toggle_agent"] == "agent_0"

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Verify RedVictim state was restored
        restored_victim = env2.grid.get(*red_victim_pos)
        assert isinstance(restored_victim, RedVictim)
        assert restored_victim.toggle_countdown == 15
        assert restored_victim.first_toggle_agent == "agent_0"

    def test_agent_with_medkit_roundtrip(self, env):
        """Give agent a MedKit in inventory, roundtrip, verify inventory restored."""
        env.reset(seed=42)

        # Give agent 0 a MedKit in inventory
        agent = env.env_agents[0]
        medkit = MedKit(state=0)
        agent.inventory = [medkit]

        # Save state
        state = env.get_state()

        # Verify agent inventory is in state
        assert len(state["agents"][0]["inventory"]) == 1
        assert state["agents"][0]["inventory"][0]["object_id"] == "medkit"

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Verify inventory was restored
        restored_agent = env2.env_agents[0]
        assert len(restored_agent.inventory) == 1
        assert isinstance(restored_agent.inventory[0], MedKit)

    def test_termination_preserved_after_restore(self, env):
        """Verify that termination state is preserved after roundtrip.

        S&R terminates when all victims are rescued (removed from grid).
        This tests the S&R-specific termination logic in SearchRescueEnv.get_terminateds_truncateds().
        """
        env.reset(seed=42)

        # Remove all victims from grid (simulating they were rescued)
        # Find and remove GreenVictim, YellowVictim, and any RedVictim
        from cogrid.envs.search_rescue.search_rescue_grid_objects import (
            GreenVictim,
            YellowVictim,
        )

        victims_to_remove = []
        for row in range(env.grid.height):
            for col in range(env.grid.width):
                obj = env.grid.get(row, col)
                if isinstance(obj, (GreenVictim, YellowVictim, RedVictim)):
                    victims_to_remove.append((row, col))

        for row, col in victims_to_remove:
            env.grid.set(row, col, None)

        # Take a step to trigger termination check
        # All victims removed means environment should terminate
        obs, rewards, terminateds, truncateds, infos = env.step(
            {0: Actions.Noop, 1: Actions.Noop}
        )

        # Verify all agents have terminated
        assert all(terminateds.values()), "Expected all agents to terminate after all victims rescued"

        # Save state after termination
        state = env.get_state()

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Verify terminated state was restored for all agents
        for agent_id, agent in env2.env_agents.items():
            assert agent.terminated, f"Agent {agent_id} should still be terminated after restore"

    def test_rng_state_preservation_sr(self, env):
        """Verify RNG state is preserved across roundtrip for reproducibility.

        After restoring from saved state, environment should produce the same
        random sequences as the original (important for MCTS and other planners).
        """
        env.reset(seed=42)

        # Take some steps to advance the state
        actions = {0: Actions.MoveRight, 1: Actions.MoveDown}
        for _ in range(5):
            env.step(actions)

        # Save state
        state = env.get_state()

        # Verify RNG state is captured
        assert "rng_state" in state
        assert state["rng_state"] is not None

        # Continue taking steps in env1 and record observations
        env1_observations = []
        for _ in range(5):
            obs, _, _, _, _ = env.step(actions)
            env1_observations.append(obs[0].copy())

        # Restore to new environment
        env2 = registry.make("SearchRescue-Test-V0")
        env2.set_state(state)

        # Take same steps in env2 and record observations
        env2_observations = []
        for _ in range(5):
            obs, _, _, _, _ = env2.step(actions)
            env2_observations.append(obs[0].copy())

        # Observations should match exactly (deterministic behavior from same RNG state)
        for i, (obs1, obs2) in enumerate(zip(env1_observations, env2_observations)):
            np.testing.assert_array_equal(
                obs1, obs2, err_msg=f"Observation mismatch at step {i+1} after restore"
            )
