import sys
import os
import time
import hashlib
import gc
import sqlite3
from typing import Dict, Any
from datetime import datetime
from cpu.enhanced_cpu import EnhancedCPU, VirtualCPU, CPUGroupType
from disk_storage import DiskStorageManager

# Project root path setup
project_root = os.path.dirname(os.path.abspath(__file__))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# Initialize storage systems
disk_storage = DiskStorageManager(storage_path="storage/mining", compression="lzma", max_size_gb=10.0)

def sha256d(header: bytes) -> bytes:
    """Perform Bitcoin's double SHA256"""
    return hashlib.sha256(hashlib.sha256(header).digest()).digest()

def process_batch_virtual(cpu_id: int, start_nonce: int, end_nonce: int, config: 'MiningConfig') -> Dict[str, Any]:
    """Process a batch using virtual CPU"""
    # Initialize virtual CPU with storage reference
    vcpu = VirtualCPU(
        cpu_id=cpu_id,
        group_type=CPUGroupType.COMPUTATION
    )
    # Share our storage connection
    vcpu.storage = config.storage
    
    results = {
        'hashes': 0,
        'blocks_found': 0,
        'best_hash': None
    }
    
    # Process nonces in small chunks to control memory
    chunk_size = 1000
    for nonce in range(start_nonce, end_nonce, chunk_size):
        chunk_end = min(nonce + chunk_size, end_nonce)
        
        # Store chunk data to disk
        chunk_data = {
            'nonces': list(range(nonce, chunk_end)),
            'headers': [f"block_header_{n}".encode() for n in range(nonce, chunk_end)]
        }
        
        # Use disk storage's compression
        disk_storage._compress_data(str(chunk_data).encode())
        
        # Process chunk using virtual CPU
        vcpu.compute_at_electron_speed(chunk_end - nonce)
        results['hashes'] += chunk_end - nonce
        
        # Force cleanup after each chunk
        gc.collect()
        
    return results

def init_storage():
    """Initialize storage with clean database"""
    if not os.path.exists("data"):
        os.makedirs("data")
    
    # Clean up any existing database
    if os.path.exists("data/cpu_storage.db"):
        os.remove("data/cpu_storage.db")
    
    # Create new database connection
    conn = sqlite3.connect("data/cpu_storage.db")
    cursor = conn.cursor()
    
    # Create all necessary tables
    cursor.execute("""
        CREATE TABLE cpu_system_config (
            id INTEGER PRIMARY KEY,
            total_cpus INTEGER,
            cores_per_cpu INTEGER,
            threads_per_core INTEGER,
            init_batch_size INTEGER,
            batch_size INTEGER,
            initialized INTEGER DEFAULT 0,
            last_updated INTEGER DEFAULT 0
        )
    """)
    
    cursor.execute("""
        CREATE TABLE cpu_config (
            cpu_id INTEGER PRIMARY KEY,
            gate_delay REAL,
            switch_freq REAL,
            drift_speed REAL,
            clock_freq REAL,
            group_type INTEGER,
            core_count INTEGER,
            thread_count INTEGER,
            initialized INTEGER DEFAULT 0,
            last_updated INTEGER DEFAULT 0
        )
    """)
    
    cursor.execute("""
        CREATE TABLE core_state (
            core_id INTEGER,
            cpu_id INTEGER,
            frequency REAL,
            gate_delay REAL,
            drift_speed REAL,
            PRIMARY KEY (core_id, cpu_id)
        )
    """)
    
    cursor.execute("""
        CREATE TABLE thread_state (
            thread_id INTEGER,
            core_id INTEGER,
            cpu_id INTEGER,
            busy INTEGER DEFAULT 0,
            last_op TEXT,
            PRIMARY KEY (thread_id, core_id, cpu_id)
        )
    """)
    
    cursor.execute("""
        CREATE TABLE compute_state (
            cpu_id INTEGER,
            cycles INTEGER,
            frequency REAL,
            latency REAL,
            throughput REAL,
            compute_time REAL,
            PRIMARY KEY (cpu_id)
        )
    """)
    
    conn.commit()
    return conn, cursor

class StorageConnection:
    """Custom storage connection that doesn't auto-create tables"""
    def __init__(self, db_path):
        self.db_path = db_path
        self._connect()
    
    def _connect(self):
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        
    def close(self):
        if self.conn:
            self.conn.close()

def init_virtual_cpu_system(total_cpus=10, cores_per_cpu=10, threads_per_core=100, batch_size=5):
    """Initialize the virtual CPU system with clean storage"""
    global storage
    
    print("Initializing virtual CPU system...")
    print("Cleaning up old database...")
    
    # Clean up old database
    try:
        if os.path.exists("data/cpu_storage.db"):
            print("Removing old database...")
            os.remove("data/cpu_storage.db")
            print("Old database removed")
    except Exception as e:
        print(f"Error removing database: {e}")
    
    # Make sure data directory exists
    if not os.path.exists("data"):
        print("Creating data directory...")
        os.makedirs("data")
        print("Data directory created")
    
    # Create and initialize storage
    storage = LocalStorageManager(db_path="data/cpu_storage.db")
    
    # Reset all tables
    storage.cursor.execute("DROP TABLE IF EXISTS cpu_system_config")
    storage.cursor.execute("DROP TABLE IF EXISTS cpu_config")
    storage.cursor.execute("DROP TABLE IF EXISTS core_state")
    storage.cursor.execute("DROP TABLE IF EXISTS thread_state")
    storage.cursor.execute("DROP TABLE IF EXISTS compute_state")
    
    # Create tables
    storage.cursor.execute("""
        CREATE TABLE cpu_system_config (
            id INTEGER PRIMARY KEY,
            total_cpus INTEGER,
            cores_per_cpu INTEGER,
            threads_per_core INTEGER,
            init_batch_size INTEGER,
            batch_size INTEGER,
            initialized INTEGER DEFAULT 0,
            last_updated INTEGER DEFAULT 0
        )
    """)
    
    storage.cursor.execute("""
        CREATE TABLE cpu_config (
            cpu_id INTEGER PRIMARY KEY,
            gate_delay REAL,
            switch_freq REAL,
            drift_speed REAL,
            clock_freq REAL,
            group_type INTEGER,
            core_count INTEGER,
            thread_count INTEGER,
            initialized INTEGER DEFAULT 0,
            last_updated INTEGER DEFAULT 0
        )
    """)
    
    storage.cursor.execute("""
        CREATE TABLE core_state (
            core_id INTEGER,
            cpu_id INTEGER,
            frequency REAL,
            gate_delay REAL,
            drift_speed REAL,
            PRIMARY KEY (core_id, cpu_id)
        )
    """)
    
    storage.cursor.execute("""
        CREATE TABLE thread_state (
            thread_id INTEGER,
            core_id INTEGER,
            cpu_id INTEGER,
            busy INTEGER DEFAULT 0,
            last_op TEXT,
            PRIMARY KEY (thread_id, core_id, cpu_id)
        )
    """)
    
    storage.cursor.execute("""
        CREATE TABLE compute_state (
            cpu_id INTEGER,
            cycles INTEGER,
            frequency REAL,
            latency REAL,
            throughput REAL,
            compute_time REAL,
            PRIMARY KEY (cpu_id)
        )
    """)
    
    # Initialize system configuration
    storage.cursor.execute("""
        INSERT INTO cpu_system_config 
        (id, total_cpus, cores_per_cpu, threads_per_core, init_batch_size, batch_size, initialized)
        VALUES (1, ?, ?, ?, ?, ?, 1)
    """, (total_cpus, cores_per_cpu, threads_per_core, batch_size, batch_size))
    
    storage.conn.commit()
    
    return storage

# Global storage reference
storage = None

class MiningConfig:
    """Configuration for mining parameters"""
    def __init__(self):
        global storage
        
        # Virtual CPU Configuration - reduced for testing
        self.total_virtual_cpus = 10  # Start with fewer CPUs for testing
        self.cores_per_cpu = 10  # Cores per CPU
        self.threads_per_core = 100  # Threads per core
        self.batch_size = 5  # Process in smaller batches for testing
        
        # Mining Parameters
        self.test_duration = 120  # seconds
        self.log_interval = 1.0
        self.nonce_range = 0xFFFFFF  # Smaller range for testing
        
        # Initialize storage and virtual CPU system
        storage = init_virtual_cpu_system(
            total_cpus=self.total_virtual_cpus,
            cores_per_cpu=self.cores_per_cpu,
            threads_per_core=self.threads_per_core,
            batch_size=self.batch_size
        )
        self.storage = storage        # Create necessary tables with updated schema
        # Create compute state table
        self.storage.cursor.execute("""
            CREATE TABLE IF NOT EXISTS compute_state (
                cpu_id INTEGER,
                cycles INTEGER,
                frequency REAL,
                latency REAL,
                throughput REAL,
                compute_time REAL,
                PRIMARY KEY (cpu_id)
            )
        """)
        
        # Create core state table
        self.storage.cursor.execute("""
            CREATE TABLE IF NOT EXISTS core_state (
                core_id INTEGER,
                cpu_id INTEGER,
                frequency REAL,
                gate_delay REAL,
                drift_speed REAL,
                PRIMARY KEY (core_id, cpu_id)
            )
        """)
        
        # Create thread state table
        self.storage.cursor.execute("""
            CREATE TABLE IF NOT EXISTS thread_state (
                thread_id INTEGER,
                core_id INTEGER,
                cpu_id INTEGER,
                busy INTEGER DEFAULT 0,
                last_op TEXT,
                PRIMARY KEY (thread_id, core_id, cpu_id)
            )
        """)
        
        self.storage.cursor.execute("""
            CREATE TABLE cpu_system_config (
                id INTEGER PRIMARY KEY,
                total_cpus INTEGER,
                cores_per_cpu INTEGER,
                threads_per_core INTEGER,
                init_batch_size INTEGER,
                batch_size INTEGER,
                initialized INTEGER DEFAULT 0,
                last_updated INTEGER DEFAULT 0
            )
        """)
        
        self.storage.cursor.execute("""
            CREATE TABLE cpu_config (
                cpu_id INTEGER PRIMARY KEY,
                gate_delay REAL,
                switch_freq REAL,
                drift_speed REAL,
                clock_freq REAL,
                group_type INTEGER,
                core_count INTEGER,
                thread_count INTEGER,
                initialized INTEGER DEFAULT 0,
                last_updated INTEGER DEFAULT 0
            )
        """)
        
        self.storage.cursor.execute("""
            CREATE TABLE core_state (
                core_id INTEGER,
                cpu_id INTEGER,
                frequency REAL,
                gate_delay REAL,
                drift_speed REAL,
                PRIMARY KEY (core_id, cpu_id)
            )
        """)
        
        self.storage.cursor.execute("""
            CREATE TABLE thread_state (
                thread_id INTEGER,
                core_id INTEGER,
                cpu_id INTEGER,
                busy INTEGER DEFAULT 0,
                last_op TEXT,
                PRIMARY KEY (thread_id, core_id, cpu_id)
            )
        """)
        
        # Insert or update system configuration
        self.storage.cursor.execute("""
            INSERT OR REPLACE INTO cpu_system_config 
            (id, total_cpus, cores_per_cpu, threads_per_core, init_batch_size, batch_size, initialized)
            VALUES (1, ?, ?, ?, ?, ?, 1)
        """, (
            self.total_virtual_cpus,
            self.cores_per_cpu,
            self.threads_per_core,
            self.batch_size,  # Use same value for both batch sizes initially
            self.batch_size
        ))
        
        self.storage.conn.commit()
        
        # Initialize virtual CPU system
        print("Initializing Virtual CPU System...")
        VirtualCPU.initialize_system(
            total_cpus=self.total_virtual_cpus,
            cores_per_cpu=self.cores_per_cpu,
            threads_per_core=self.threads_per_core,
            batch_size=self.batch_size
        )
        print("Virtual CPU System Initialized")

class MiningStats:
    """Track mining statistics using disk storage"""
    def __init__(self):
        self.start_time = time.time()
        self.last_update = time.time()
        self.stats_path = "storage/mining/stats"
        os.makedirs(self.stats_path, exist_ok=True)
        
    def update_stats(self, thread_results: Dict[str, Any]):
        """Update stats using disk storage"""
        timestamp = datetime.now().isoformat()
        disk_storage.store_batch(
            f"stats_{timestamp}",
            [{
                'timestamp': timestamp,
                'results': thread_results
            }]
        )

def main():
    config = MiningConfig()
    stats = MiningStats()
    
    print("Starting mining test with virtual CPUs...")
    
    # Calculate work distribution
    total_nonces = config.nonce_range
    nonces_per_cpu = (total_nonces + config.total_virtual_cpus - 1) // config.total_virtual_cpus
    
    start_time = time.time()
    active_cpus = []
    
    try:
        # Initialize virtual CPUs in batches
        for cpu_batch in range(0, config.total_virtual_cpus, config.batch_size):
            batch_end = min(cpu_batch + config.batch_size, config.total_virtual_cpus)
            
            for cpu_id in range(cpu_batch, batch_end):
                start_nonce = cpu_id * nonces_per_cpu
                end_nonce = min(start_nonce + nonces_per_cpu, total_nonces)
                
                # Process batch with virtual CPU
                results = process_batch_virtual(cpu_id, start_nonce, end_nonce, config)
                stats.update_stats(results)
                
                # Force cleanup after each CPU
                gc.collect()
            
            # Periodic cleanup
            if time.time() - start_time > config.test_duration:
                break
                
    except KeyboardInterrupt:
        print("\nStopping mining test...")
    finally:
        # Manual cleanup
        gc.collect()
        # Clear any temporary files
        for filename in os.listdir("storage/mining"):
            filepath = os.path.join("storage/mining", filename)
            try:
                if os.path.isfile(filepath):
                    os.unlink(filepath)
            except Exception as e:
                print(f'Error: {e}')

if __name__ == "__main__":
    main()