import time
import random
import hashlib
import multiprocessing
from concurrent.futures import ProcessPoolExecutor, Future
from typing import List, Optional, Tuple
from prophetdiceskills.core import StakeRNG
from prophetdiceskills import gpu_miner

def worker_cpu_collision(args: Tuple) -> Optional[str]:
    prefix_str, suffix_range, alphabet_str, client_seed, num_targets, target_outcomes, start_nonce, public_hash, suffix_len = args
    for offset in suffix_range:
        temp = offset
        seed_str = ""
        for _ in range(suffix_len): 
            seed_str = alphabet_str[temp % 16] + seed_str
            temp //= 16
        test_seed = prefix_str + seed_str
        
        # Immediate Fast-Path: Does the generated full seed match the public commitment?
        if public_hash:
             if hashlib.sha256(test_seed.encode()).hexdigest() == public_hash:
                 return test_seed
                
        # Look for a functional collision (a seed that produces the identical history)
        rng = StakeRNG(test_seed, client_seed)
        all_matched = True
        for t_idx, target in enumerate(target_outcomes):
            roll = rng.get_dice_roll(start_nonce + t_idx)
            if abs(roll - target) >= 0.001:
                all_matched = False
                break
        if all_matched:
            return test_seed
    return None

def worker_gpu_collision(args: Tuple) -> Optional[str]:
    """Wraps the Numba device execution so it runs cleanly isolated in a background process"""
    if not gpu_miner.HAS_CUDA: return None
    prefix_str, client_seed, num_targets, target_outcomes, start_nonce, public_hash, target_length = args
    
    try:
        import numpy as np
        from numba import cuda
        
        alphabet_str = "0123456789abcdef"
        alphabet_dev = cuda.to_device(np.frombuffer(alphabet_str.encode('ascii'), dtype=np.uint8))
        msg_base_str = f"{client_seed}:"
        msg_base_bytes = msg_base_str.encode('utf-8')
        client_msg_base_len = len(msg_base_bytes)
        
        client_msg_base_dev = cuda.to_device(np.frombuffer(msg_base_bytes, dtype=np.uint8))
        target_outcomes_dev = cuda.to_device(np.array(target_outcomes, dtype=np.float64))
        k_dev = cuda.to_device(gpu_miner.K)
        
        result_found_host = np.zeros(128, dtype=np.int32)
        result_len_host = np.zeros(1, dtype=np.int32)
        result_seed_idx_host = np.zeros(1, dtype=np.uint64)
        result_found = cuda.to_device(result_found_host)
        result_len = cuda.to_device(result_len_host)
        result_seed_idx = cuda.to_device(result_seed_idx_host)
        
        prefix_len = len(prefix_str)
        suffix_len = target_length - prefix_len
        total_perms = 16 ** suffix_len
        
        threads_per_block = 128
        chunk_size = 1024 * 512
        if chunk_size % threads_per_block != 0:
             chunk_size += threads_per_block - (chunk_size % threads_per_block)
             
        workspace_size = chunk_size * 64
        workspace_dev = cuda.device_array(workspace_size, dtype=np.uint32)
        
        prefix_bytes = prefix_str.encode('ascii')
        prefix_dev = cuda.to_device(np.frombuffer(prefix_bytes, dtype=np.uint8))
        stream = cuda.stream()
        
        for offset in range(0, total_perms, chunk_size):
            current_batch = min(chunk_size, total_perms - offset)
            blocks = (current_batch + threads_per_block - 1) // threads_per_block
            
            gpu_miner.brute_force_kernel_beatstake[blocks, threads_per_block, stream](
                k_dev, client_msg_base_dev, client_msg_base_len, target_outcomes_dev, num_targets,
                result_found, result_len, result_seed_idx,
                suffix_len, alphabet_dev, offset, workspace_dev, prefix_dev, prefix_len, start_nonce
            )
            result_found.copy_to_host(result_found_host, stream=stream)
            stream.synchronize()
            
            if result_found_host[0] == 1:
                result_seed_idx.copy_to_host(result_seed_idx_host, stream=stream)
                stream.synchronize()
                seed_idx = int(result_seed_idx_host[0])
                seed_str = ""
                temp = seed_idx
                for _ in range(suffix_len):
                    seed_str = alphabet_str[temp % 16] + seed_str
                    temp //= 16
                final_seed = prefix_str + seed_str
                
                # Double check with public hash if provided
                if public_hash:
                     if hashlib.sha256(final_seed.encode()).hexdigest() == public_hash:
                         return final_seed
                else:
                    return final_seed
                    
        return None
    except Exception as e:
        print(f"Background GPU Error: {e}")
        return None

class BackgroundSeedMiner:
    """
    Constantly runs in the background attempting to crack the active server seed
    either by timestamp reversing or collision farming.
    """
    def __init__(self, cores: int = None):
        self.cores = cores or max(1, multiprocessing.cpu_count() - 1)
        self.active = False
        self.pool = None
        self.current_futures = []
        self.cracked_seed = None
        self.mining_active = False

    def start(self):
        self.active = True
        self.pool = ProcessPoolExecutor(max_workers=self.cores)

    def stop(self):
        self.active = False
        if self.pool:
            self.pool.shutdown(wait=False, cancel_futures=True)
            self.pool = None
            
    def get_cracked_seed(self) -> Optional[str]:
        # Non-blocking poll of background tasks
        if not self.current_futures:
            return self.cracked_seed
            
        still_running = []
        for f in self.current_futures:
            if f.done():
                result = f.result()
                if result:
                    self.cracked_seed = result
                    self.stop() # Stop mining once found
                    return result
            else:
                still_running.append(f)
                
        self.current_futures = still_running
        if not still_running:
            self.mining_active = False
        return self.cracked_seed

    def farm_collision_async(self, client_seed: str, target_outcomes: List[float], start_nonce: int = 1, target_length: int = 64, public_hash: Optional[str] = None, known_prefix: Optional[str] = None):
        """
        Dispatches a cracking job to background workers if no job is currently running.
        If a known_prefix is provided, it only searches the suffix permutations for that prefix.
        """
        if self.cracked_seed or self.mining_active:
            return # Already found or already searching
            
        self.mining_active = True
            
        if not self.pool:
            self.start()
            
        if len(self.current_futures) >= self.cores * 2:
            return # Queue full
            
        alphabet_str = "0123456789abcdef"
        
        if known_prefix:
            prefix_len = len(known_prefix)
            suffix_len = target_length - prefix_len
            if suffix_len <= 0: return []
            prefix_str = known_prefix
        else:
            suffix_len = 5 # Faster iterations for continuous background
            prefix_len = target_length - suffix_len
            prefix_str = "".join(random.choices(alphabet_str, k=prefix_len))
            
        # Dispatch Logic
        if gpu_miner.HAS_CUDA:
            # GPU handles the entire block space in one device kernel array
            chunk = (prefix_str, client_seed, len(target_outcomes), target_outcomes, start_nonce, public_hash, target_length)
            self.current_futures.append(self.pool.submit(worker_gpu_collision, chunk))
        else:
            # CPU strips the workload across all standard logic cores
            total_perms = 16 ** suffix_len
            chunk_size = max(1, total_perms // self.cores)
            work_chunks = []
            for i in range(self.cores):
                start_off = i * chunk_size
                end_off = (i + 1) * chunk_size if i < self.cores - 1 else total_perms
                work_chunks.append((prefix_str, range(start_off, end_off), alphabet_str, client_seed, len(target_outcomes), target_outcomes, start_nonce, public_hash, suffix_len))
                
            for chunk in work_chunks:
                self.current_futures.append(self.pool.submit(worker_cpu_collision, chunk))

