import os
import sys
import unittest
import random
import json
import hashlib
from pathlib import Path

# Add current directory to path
sys.path.append(os.getcwd())

# Set SDL to use dummy video driver
os.environ["SDL_VIDEODRIVER"] = "dummy"
os.environ["SDL_AUDIODRIVER"] = "dummy"
os.environ["MICE_DISABLE_JOYSTICK"] = "1"

from rats import MiceMaze
from engine.state_machine import GameState
from engine.sdl2 import GameWindow

def mock_init_audio(self):
    self.music_enabled = False
    self.audio_devs = {"base": 0, "effects": 0, "music": 0}
    self.sound_volume = 0
    self.music_volume = 0

class LoopParityTester(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        GameWindow.show_intro = lambda *args, **kwargs: None
        GameWindow.show_loading_screen = lambda *args, **kwargs: None
        GameWindow._init_audio_system = mock_init_audio
        GameWindow.play_sound = lambda *args, **kwargs: None
        GameWindow.stop_sound = lambda *args, **kwargs: None
        
    def setUp(self):
        # We need absolute determinism
        random.seed(12345)
        self.game = MiceMaze("assets/Rat/level.dat", level_index=0)
        
        # Reset and restart to clear initialization entropy
        random.seed(12345)
        # Monkeypatch UUID to be deterministic
        import uuid
        self.uuid_counter = 0
        def mock_uuid4():
            self.uuid_counter += 1
            return self.uuid_counter
        uuid.uuid4 = mock_uuid4
        
        self.game.start_game()
        self.game.state_machine.transition_to(GameState.PLAYING)
        
    def get_full_snapshot(self):
        """Captures extremely detailed state of all units."""
        snapshot = {
            "points": self.game.points,
            "units": []
        }
        # Sort by ID for stability
        sorted_units = sorted(self.game.units.items(), key=lambda x: int(x[0]))
        
        for uid, u in sorted_units:
            u_data = {
                "id": uid,
                "type": u.__class__.__name__,
                "pos": list(u.position),
                "pos_before": list(u.position_before),
                "partial": float(u.partial_move),
                "age": int(u.age)
            }
            # Optional attributes that affect logic
            if hasattr(u, "pregnant"): u_data["pregnant"] = int(u.pregnant)
            if hasattr(u, "babies"): u_data["babies"] = int(u.babies)
            if hasattr(u, "gassed"): u_data["gassed"] = int(u.gassed)
            if hasattr(u, "direction"): u_data["dir"] = u.direction
            
            snapshot["units"].append(u_data)
            
        # Add a hash of the total unit count and positions for quick check
        flat_state = str(snapshot).encode('utf-8')
        snapshot["hash"] = hashlib.md5(flat_state).hexdigest()
        return snapshot

    def test_record_or_verify(self):
        steps = 100
        parity_file = Path("tests/loop_parity_master.json")
        
        states = []
        print(f"Running simulation for {steps} steps...")
        
        for i in range(steps):
            self.game.update_maze()
            states.append(self.get_full_snapshot())
            
        if not parity_file.exists() or os.environ.get("RECORD_PARITY"):
            with open(parity_file, "w") as f:
                json.dump(states, f, indent=2)
            print(f"RECORDED master state to {parity_file}")
        else:
            with open(parity_file, "r") as f:
                master_states = json.load(f)
            
            self.assertEqual(len(states), len(master_states))
            
            for i, (curr, master) in enumerate(zip(states, master_states)):
                if curr["hash"] != master["hash"]:
                    # Detailed comparison on failure
                    self.assertEqual(curr["points"], master["points"], f"Points mismatch at step {i}")
                    self.assertEqual(len(curr["units"]), len(master["units"]), f"Unit count mismatch at step {i}")
                    
                    for u_idx, (u_curr, u_master) in enumerate(zip(curr["units"], master["units"])):
                        self.assertEqual(u_curr, u_master, f"Unit {u_idx} mismatch at step {i}")
            
            print("PARITY VERIFIED: Optimization is functionally identical.")

if __name__ == "__main__":
    unittest.main()