"""
JSON Storage Manager Implementation
Uses JSON files for transparent storage and real-time monitoring
"""

import json
import os
import time
import threading
from typing import Dict, Any, Optional, Union, List

from pathlib import Path
import atexit

class JSONStorageManager:
    """
    JSON-based storage implementation for transparent mining operations.
    
    Features:
    - Live monitoring through JSON files
    - Separate files for different data types
    - Real-time stats updates
    - Thread-safe operations
    """
    
    def __init__(self, base_path: str = "data"):
        self.base_path = Path(base_path)
        self.locks = {
            'cpu_stats': threading.Lock(),
            'mining_results': threading.Lock(),
            'system_config': threading.Lock(),
            'hash_stats': threading.Lock()
        }
        self.files = {
            'cpu_stats': self.base_path / 'cpu_stats.json',
            'mining_results': self.base_path / 'mining_results.json',
            'system_config': self.base_path / 'system_config.json',
            'hash_stats': self.base_path / 'hash_stats.json'
        }
        
        # Initialize storage
        self._init_storage()
        
        # Register cleanup on exit
        atexit.register(self.cleanup)
        
    def _init_storage(self):
        """Initialize JSON storage files"""
        os.makedirs(self.base_path, exist_ok=True)
        
        # Initialize CPU stats
        if not self.files['cpu_stats'].exists():
            self._save_json(self.files['cpu_stats'], {
                'cpus': {},
                'last_updated': time.time()
            })
        
        # Initialize mining results
        if not self.files['mining_results'].exists():
            self._save_json(self.files['mining_results'], {
                'results': [],
                'best_hash': None,
                'best_nonce': None
            })
            
        # Initialize system config
        if not self.files['system_config'].exists():
            self._save_json(self.files['system_config'], {
                'total_cpus': 1000,
                'cores_per_cpu': 1000,
                'threads_per_core': 1000,
                'initialized': False
            })
            
        # Initialize hash stats
        if not self.files['hash_stats'].exists():
            self._save_json(self.files['hash_stats'], {
                'total_hashes': 0,
                'hashes_per_cpu': {},
                'hashes_per_core': {},
                'start_time': time.time()
            })
            
    def _load_json(self, file_path: Path) -> dict:
        """Load JSON file with lock protection and retry logic"""
        max_retries = 3
        for attempt in range(max_retries):
            with self.locks[file_path.stem]:
                try:
                    with open(file_path, 'r') as f:
                        data = f.read()
                        if not data:  # Empty file
                            return {}
                        return json.loads(data)
                except (json.JSONDecodeError, FileNotFoundError):
                    if attempt == max_retries - 1:
                        return {}
                    time.sleep(0.1)
                
    def _save_json(self, file_path: Path, data: dict):
        """Save JSON file with lock protection"""
        with self.locks[file_path.stem]:
            try:
                # Direct write for better Windows compatibility
                with open(file_path, 'w') as f:
                    json.dump(data, f, indent=2)
            except Exception as e:
                print(f"Error saving JSON to {file_path}: {e}")
                raise
                
    def update_cpu_stats(self, cpu_id: int, stats: dict):
        """Update stats for a specific CPU"""
        data = self._load_json(self.files['cpu_stats'])
        data['cpus'][str(cpu_id)] = stats
        data['last_updated'] = time.time()
        self._save_json(self.files['cpu_stats'], data)
        
    def record_hash(self, cpu_id: int, core_id: int, thread_id: int, hash_count: int):
        """Record hashes performed by CPU/core/thread"""
        print(f"[DEBUG-storage] record_hash called: cpu_id={cpu_id}, core_id={core_id}, thread_id={thread_id}, hash_count={hash_count}")
        data = self._load_json(self.files["hash_stats"])
        
        # Update total hashes
        if 'total_hashes' not in data:
            data['total_hashes'] = 0
        data["total_hashes"] += hash_count
        
        # Update per-CPU hashes
        cpu_key = str(cpu_id)
        if 'hashes_per_cpu' not in data:
            data['hashes_per_cpu'] = {}
        if cpu_key not in data["hashes_per_cpu"]:
            data["hashes_per_cpu"][cpu_key] = 0
        data["hashes_per_cpu"][cpu_key] += hash_count
        
        # Update per-core hashes
        core_key = f"{cpu_id}:{core_id}"
        if 'hashes_per_core' not in data:
            data['hashes_per_core'] = {}
        if core_key not in data["hashes_per_core"]:
            data["hashes_per_core"][core_key] = 0
        data["hashes_per_core"][core_key] += hash_count

        # Update per-thread hashes
        thread_key = f"{cpu_id}:{core_id}:{thread_id}"
        if 'hashes_per_thread' not in data:
            data['hashes_per_thread'] = {}
        if thread_key not in data["hashes_per_thread"]:
            data["hashes_per_thread"][thread_key] = 0
        data["hashes_per_thread"][thread_key] += hash_count
        
        self._save_json(self.files["hash_stats"], data)
        print(f"[DEBUG-storage] hash_stats after record_hash: {data}")
        
    def add_mining_result(self, result: dict):
        """Add a mining result"""
        print(f"[DEBUG-storage] add_mining_result called: result={result}")
        data = self._load_json(self.files["mining_results"])
        data["results"].append({
            **result,
            "timestamp": time.time()
        })
        
        # Update best hash if needed
        if not data["best_hash"] or result["hash"] < data["best_hash"]:
            data["best_hash"] = result["hash"]
            data["best_nonce"] = result["nonce"]
            
        self._save_json(self.files["mining_results"], data)
        print(f"[DEBUG-storage] mining_results after add_mining_result: {data}")
        
    def get_system_config(self) -> dict:
        """Get current system configuration"""
        return self._load_json(self.files["system_config"])
        
    def update_system_config(self, config: dict):
        """Update system configuration"""
        current = self.get_system_config()
        current.update(config)
        self._save_json(self.files["system_config"], current)
        
    def get_mining_stats(self) -> dict:
        """Get current mining statistics"""
        cpu_stats = self._load_json(self.files["cpu_stats"])
        hash_stats = self._load_json(self.files["hash_stats"])
        mining_results = self._load_json(self.files["mining_results"])
        
        # Ensure all expected keys are present in hash_stats
        hash_stats.setdefault('total_hashes', 0)
        hash_stats.setdefault('hashes_per_cpu', {})
        hash_stats.setdefault('hashes_per_core', {})
        hash_stats.setdefault('hashes_per_thread', {})
        hash_stats.setdefault('start_time', time.time())

        print(f"[DEBUG-storage] get_mining_stats returning: hash_stats={hash_stats}, best_result={mining_results}")
        return {
            "cpu_stats": cpu_stats,
            "hash_stats": {
                "total_hashes": hash_stats["total_hashes"],
                "hashes_per_cpu": hash_stats["hashes_per_cpu"],
                "hashes_per_core": hash_stats["hashes_per_core"],
                "hashes_per_thread": hash_stats["hashes_per_thread"],
                "start_time": hash_stats["start_time"]
            },
            "best_result": {
                "hash": mining_results.get("best_hash"),
                "nonce": mining_results.get("best_nonce")
            },
            "total_hashes": hash_stats["total_hashes"],
            "runtime": time.time() - hash_stats["start_time"]
        }
        
    def cleanup(self):
        """Cleanup on exit"""
        # Save final stats
        stats = self.get_mining_stats()
        final_stats_file = self.base_path / 'final_stats.json'
        with open(final_stats_file, 'w') as f:
            json.dump(stats, f, indent=2)

    def update_mining_stats(self, stats: dict):
        """Update overall mining statistics"""
        data = self._load_json(self.files["hash_stats"])
        data["total_hashes"] = stats.get("total_hashes", data.get("total_hashes", 0))
        data["total_blocks_found"] = stats.get("total_blocks_found", data.get("total_blocks_found", 0))
        
        mining_results = self._load_json(self.files["mining_results"])
        if stats.get("best_hash") and (not mining_results.get("best_hash") or stats["best_hash"] < mining_results["best_hash"]):
            mining_results["best_hash"] = stats["best_hash"]
            mining_results["best_nonce"] = stats["best_nonce"]
        self._save_json(self.files["mining_results"], mining_results)
        
        self._save_json(self.files["hash_stats"], data)

