"""
JSON-based state manager for CPU state persistence
"""
import json
import os
import time
from pathlib import Path
from typing import Optional, Dict, Any, Tuple

class JSONStateManager:
    """Manages CPU state persistence using JSON files"""
    
    def __init__(self, base_path: str = "db/cpu"):
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)
        self.state_file = self.base_path / "cpu_states.json"
        self.states: Dict[int, Dict[str, Any]] = {}
        self._load_states()
    
    def _load_states(self):
        """Load existing states from JSON file"""
        if self.state_file.exists():
            try:
                with open(self.state_file, 'r') as f:
                    self.states = json.load(f)
            except json.JSONDecodeError:
                self.states = {}
        else:
            self.states = {}
    
    def _save_states(self):
        """Save states to JSON file"""
        with open(self.state_file, 'w') as f:
            json.dump(self.states, f, indent=2)
    
    def get_cpu_state(self, cpu_id: int) -> Tuple[Optional[dict], Optional[dict]]:
        """Get CPU state and registers from storage"""
        try:
            if str(cpu_id) in self.states:
                state = self.states[str(cpu_id)]
                return state.get('cpu_state'), state.get('registers')
            return None, None
        except Exception as e:
            print(f"Error retrieving CPU state for CPU {cpu_id}: {e}")
            return None, None
    
    def update_cpu_state(self, cpu, registers=None):
        """Update CPU state and registers in storage"""
        try:
            state = {
                'cpu_state': cpu.to_json(),
                'registers': registers.to_json() if registers else {},
                'last_updated': time.time_ns()
            }
            self.states[str(cpu.cpu_id)] = state
            self._save_states()
        except Exception as e:
            print(f"Error updating CPU state: {e}")
            raise
    
    def clear_state(self, cpu_id: int):
        """Clear CPU state from storage"""
        if str(cpu_id) in self.states:
            del self.states[str(cpu_id)]
            self._save_states()
