"""
EM Module - Session Integration

Provides session save/load hooks specifically for EM module state.
"""

from typing import Optional, Dict, Any
from session_integration import SessionIntegration


class EMSessionIntegration:
    """
    Extends SessionIntegration with EM-specific state handling.
    """
    
    def __init__(self, session_integration: SessionIntegration):
        self.session_integration = session_integration
    
    def capture_em_state(self, trame_state: Any) -> Dict[str, Any]:
        """Capture EM-specific state variables."""
        em_keys = [
            # Geometry
            "grid_size",
            "geometry_type",
            "geometry_rotation_x",
            "geometry_rotation_y",
            "geometry_rotation_z",
            
            # Excitation
            "excitation_type",
            "excitation_frequency",
            "excitation_location_x",
            "excitation_location_y",
            "excitation_location_z",
            
            # Simulation
            "frequency",
            "num_qubits",
            "t_final",
            "dt",
            "outer_boundary_condition",
            
            # Backend
            "backend_type",
            "selected_simulator",
            "selected_qpu",
            "aqc_enabled",
            "measurement_type",
        ]
        
        state_dict = {}
        for key in em_keys:
            if hasattr(trame_state, key):
                try:
                    value = getattr(trame_state, key)
                    # Try to serialize to ensure it's JSON-compatible
                    import json
                    json.dumps(value)
                    state_dict[key] = value
                except (TypeError, ValueError):
                    pass
        
        return state_dict
    
    def restore_em_state(self, trame_state: Any) -> None:
        """Restore EM-specific state from session."""
        if not self.session_integration.current_state:
            return
        
        for key, value in self.session_integration.current_state.state_data.items():
            if hasattr(trame_state, key):
                try:
                    setattr(trame_state, key, value)
                except Exception as e:
                    print(f"Warning: Could not restore {key}: {e}")
    
    def save_em_session(self, trame_state: Any) -> bool:
        """Save EM session with state capture."""
        state_dict = self.capture_em_state(trame_state)
        if self.session_integration.current_state:
            self.session_integration.current_state.state_data.update(state_dict)
        return self.session_integration.save_current_session(trame_state)