import secrets
import multiprocessing
import time
from datetime import datetime
from coincurve import PrivateKey
from eth_hash.auto import keccak
import pandas as pd
import os
import threading
import hashlib
import sys
import queue
import gradio as gr
from pathlib import Path
import webbrowser

# ============================================================================
# ASYNC FILE SAVER
# ============================================================================

class AsyncFileSaver:
    """Non-blocking file saving to avoid I/O bottlenecks"""

    def __init__(self):
        self.save_queue = queue.Queue()
        self.stop_event = threading.Event()
        self.worker_thread = threading.Thread(target=self._process_saves, daemon=True)
        self.worker_thread.start()

    def _process_saves(self):
        while not self.stop_event.is_set():
            try:
                data = self.save_queue.get(timeout=0.1)
                if data is None:
                    break
                address, private_key, worker_id, attempts = data
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

                with open("found_collisions.txt", "a") as f:
                    f.write(f"\n{'='*80}\n")
                    f.write(f"Timestamp: {timestamp}\n")
                    f.write(f"Worker ID: {worker_id}\n")
                    f.write(f"Attempts: {attempts:,}\n")
                    f.write(f"Address: 0x{address}\n")
                    f.write(f"Private Key: {private_key}\n")
                    f.write(f"{'='*80}\n")

                unique_filename = f"found_key_{address[:10]}_{timestamp.replace(' ', '_').replace(':', '-')}.txt"
                with open(unique_filename, "w") as f:
                    f.write(f"Address: 0x{address}\n")
                    f.write(f"Private Key: {private_key}\n")
                    f.write(f"Timestamp: {timestamp}\n")
                    f.write(f"Worker ID: {worker_id}\n")
                    f.write(f"Attempts: {attempts:,}\n")

            except queue.Empty:
                continue
            except Exception as e:
                print(f"Error saving: {e}")

    def save(self, address, private_key, worker_id, attempts):
        self.save_queue.put((address, private_key, worker_id, attempts))

    def stop(self):
        self.save_queue.put(None)
        self.stop_event.set()


# ============================================================================
# OPTIMIZED HIGH ENTROPY GENERATOR
# ============================================================================

class OptimizedHighEntropyGenerator:
    """Generate high-quality entropy from multiple sources - OPTIMIZED VERSION"""

    _static_entropy = None

    @classmethod
    def get_static_entropy(cls):
        if cls._static_entropy is None:
            entropy = b''
            entropy += str(os.getpid()).encode()
            entropy += str(os.getppid()).encode()
            entropy += str(threading.get_ident()).encode()
            entropy += str(threading.active_count()).encode()
            entropy += str(id(object())).encode()
            entropy += str(id(threading.current_thread())).encode()
            entropy += str(hash(frozenset(os.environ.items()))).encode()
            entropy += sys.version.encode()
            entropy += str(sys.maxsize).encode()
            cls._static_entropy = entropy
        return cls._static_entropy

    @staticmethod
    def get_batch_timing_jitter(batch_size):
        jitter = []
        start_base = time.perf_counter_ns()
        for i in range(min(batch_size, 256)):
            ops = 50 + (i % 100)
            start = time.perf_counter_ns()
            for j in range(ops):
                _ = j * j
            end = time.perf_counter_ns()
            jitter_value = ((end - start) ^ (start_base >> (i % 8))) & 0xFF
            jitter.append(jitter_value)

        while len(jitter) < batch_size:
            jitter.extend(jitter[:min(len(jitter), batch_size - len(jitter))])

        return bytes(jitter[:batch_size])

    @staticmethod
    def get_batch_timing_entropy(batch_size):
        samples = []
        sample_count = min(batch_size, 100)

        for i in range(sample_count):
            samples.append(time.perf_counter_ns())
            samples.append(time.monotonic_ns())
            samples.append(time.process_time_ns())
            if i % 20 == 0:
                time.sleep(0.000001 * (i % 10))

        entropy_str = b''.join(str(s).encode() for s in samples)

        while len(entropy_str) < batch_size * 32:
            entropy_str += entropy_str

        return entropy_str[:batch_size * 32]

    @classmethod
    def generate_batch_private_keys(cls, batch_size=500):
        static_entropy = cls.get_static_entropy()
        base_csprng = secrets.token_bytes(32)
        base_urandom = os.urandom(32)
        timing_jitter = cls.get_batch_timing_jitter(batch_size)
        timing_entropy = cls.get_batch_timing_entropy(batch_size)

        private_keys = []

        for i in range(batch_size):
            entropy_pool = b''
            entropy_pool += base_csprng
            entropy_pool += base_urandom
            entropy_pool += i.to_bytes(4, 'big')
            entropy_pool += timing_jitter[i:i+1] if i < len(timing_jitter) else timing_jitter[-1:]
            entropy_pool += static_entropy
            start_idx = (i * 32) % max(1, len(timing_entropy) - 32)
            entropy_pool += timing_entropy[start_idx:start_idx + 32]
            entropy_pool += os.urandom(16)
            entropy_pool += secrets.token_bytes(8)

            private_key_bytes = hashlib.sha256(entropy_pool).digest()
            private_keys.append(private_key_bytes.hex())

        return private_keys


# ============================================================================
# WORKER FUNCTION
# ============================================================================

def optimized_worker(worker_id, address_dict, total_attempts, found_queue, batch_size=500, stop_event=None):
    """Worker function with batch processing"""
    local_count = 0
    generator = OptimizedHighEntropyGenerator()

    while True:
        if stop_event and stop_event.is_set():
            break
            
        try:
            batch_keys = generator.generate_batch_private_keys(batch_size)

            for private_key_hex in batch_keys:
                private_key_bytes = bytes.fromhex(private_key_hex)
                pk = PrivateKey(private_key_bytes)
                pub = pk.public_key.format(compressed=False)[1:]
                addr_bytes = keccak(pub)[-20:]
                addr_hex = addr_bytes.hex()

                local_count += 1

                if addr_hex in address_dict:
                    found_queue.put({
                        'address': addr_hex,
                        'private_key': private_key_hex,
                        'worker_id': worker_id,
                        'attempts': local_count,
                        'timestamp': datetime.now().isoformat()
                    })
                    time.sleep(0.01)

            with total_attempts.get_lock():
                total_attempts.value += batch_size

        except Exception as e:
            print(f"Worker {worker_id} error: {e}")
            time.sleep(0.1)


# ============================================================================
# COLLISION FINDER CLASS
# ============================================================================

class CollisionFinder:
    """Manages the collision finding process"""
    
    def __init__(self):
        self.processes = []
        self.manager = None
        self.found_queue = None
        self.total_attempts = None
        self.stop_event = None
        self.monitor_thread = None
        self.saver = None
        self.is_running = False
        self.start_time = None
        
    def start_search(self, csv_file, num_workers, batch_size, test_address=None):
        """Start the collision search"""
        if self.is_running:
            return "Already running!", 0, 0, "🟢 Running", []
            
        try:
            # Load addresses
            if test_address:
                # Test mode
                address_dict = {}
                addr_clean = test_address.lower()
                if addr_clean.startswith('0x'):
                    addr_clean = addr_clean[2:]
                address_dict[addr_clean] = True
                address_count = 1
            else:
                # Load from CSV
                if csv_file is None:
                    return "❌ Please upload a CSV file first!", 0, 0, "⚪ Idle", []
                
                # Handle Gradio file object
                if hasattr(csv_file, 'name'):
                    file_path = csv_file.name
                else:
                    file_path = str(csv_file)
                
                if not os.path.exists(file_path):
                    return f"❌ File not found: {file_path}", 0, 0, "⚪ Idle", []
                
                address_df = pd.read_csv(file_path)
                address_df['address'] = address_df['address'].str.lower()
                
                address_dict = {}
                for addr in address_df['address']:
                    addr_clean = addr
                    if addr_clean.startswith('0x'):
                        addr_clean = addr_clean[2:]
                    address_dict[addr_clean] = True
                address_count = len(address_dict)
            
            # Create shared objects
            self.manager = multiprocessing.Manager()
            self.found_queue = self.manager.Queue()
            self.total_attempts = multiprocessing.Value('q', 0)
            self.stop_event = multiprocessing.Event()
            self.start_time = time.time()
            
            # Start worker processes
            self.processes = []
            for i in range(num_workers):
                p = multiprocessing.Process(
                    target=optimized_worker,
                    args=(i, address_dict, self.total_attempts, self.found_queue, batch_size, self.stop_event)
                )
                p.start()
                self.processes.append(p)
                time.sleep(0.1)
            
            # Initialize file saver
            self.saver = AsyncFileSaver()
            self.is_running = True
            
            return f"✅ Search started! Searching {address_count:,} addresses with {num_workers} workers", 0, 0, "🟢 Running", []
            
        except Exception as e:
            return f"❌ Error: {str(e)}", 0, 0, "⚪ Idle", []
    
    def stop_search(self):
        """Stop the collision search"""
        if not self.is_running:
            return "Not running!", "⚪ Idle"
            
        try:
            self.is_running = False
            if self.stop_event:
                self.stop_event.set()
            
            for p in self.processes:
                if p.is_alive():
                    p.terminate()
                    p.join(timeout=1)
            
            if self.saver:
                self.saver.stop()
            
            self.processes = []
            
            return "✅ Search stopped successfully!", "⚪ Stopped"
        except Exception as e:
            return f"❌ Error stopping: {str(e)}", "⚪ Error"
    
    def get_stats(self):
        """Get current statistics"""
        if not self.is_running:
            return 0, 0
        
        try:
            attempts = self.total_attempts.value if self.total_attempts else 0
            found_count = 0
            
            # Count found items without removing them
            temp_items = []
            while not self.found_queue.empty():
                try:
                    item = self.found_queue.get_nowait()
                    temp_items.append(item)
                except:
                    break
            
            found_count = len(temp_items)
            
            # Put items back
            for item in temp_items:
                self.found_queue.put(item)
            
            return attempts, found_count
        except:
            return 0, 0
    
    def get_found_keys(self):
        """Get list of found collisions"""
        found_keys = []
        
        if self.found_queue:
            temp_items = []
            while not self.found_queue.empty():
                try:
                    item = self.found_queue.get_nowait()
                    temp_items.append(item)
                    found_keys.append(item)
                except:
                    break
            
            # Put items back
            for item in temp_items:
                self.found_queue.put(item)
        
        return found_keys
    
    def get_update_data(self):
        """Get data for UI updates"""
        if not self.is_running:
            return 0, 0, []
        
        attempts, found = self.get_stats()
        found_keys = self.get_found_keys()
        
        # Create dataframe for found keys
        if found_keys:
            df_data = [
                [k['timestamp'], f"0x{k['address']}", k['private_key'], k['worker_id'], k['attempts']]
                for k in found_keys
            ]
        else:
            df_data = []
        
        return attempts, found, df_data


# ============================================================================
# GRADIO INTERFACE
# ============================================================================

# Initialize global finder
finder = CollisionFinder()

def create_interface():
    with gr.Blocks(title="Ethereum Address Collision Finder") as app:
        gr.Markdown("""
        # 🔑 Ethereum Address Collision Finder
        ### Search for Ethereum address collisions with high-performance multi-processing
        """)
        
        with gr.Tabs():
            # Main Search Tab
            with gr.Tab("🎯 Main Search"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 📂 Data Source")
                        csv_file = gr.File(
                            label="Upload Addresses CSV",
                            file_types=[".csv"],
                            type="filepath"
                        )
                        
                        with gr.Accordion("🔧 Advanced Settings", open=True):
                            num_workers = gr.Slider(
                                minimum=1,
                                maximum=multiprocessing.cpu_count(),
                                value=max(1, multiprocessing.cpu_count() // 2),
                                step=1,
                                label="Number of Workers"
                            )
                            batch_size = gr.Slider(
                                minimum=100,
                                maximum=2000,
                                value=500,
                                step=100,
                                label="Batch Size"
                            )
                        
                        with gr.Row():
                            start_btn = gr.Button("▶️ Start Search", variant="primary", size="lg")
                            stop_btn = gr.Button("⏹️ Stop Search", variant="stop", size="lg")
                        
                        status_text = gr.Textbox(
                            label="Status",
                            interactive=False,
                            lines=3
                        )
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### 📊 Real-time Statistics")
                        
                        with gr.Row():
                            total_attempts_display = gr.Number(
                                label="Total Attempts",
                                value=0,
                                interactive=False
                            )
                            found_count_display = gr.Number(
                                label="Collisions Found",
                                value=0,
                                interactive=False
                            )
                        
                        status_indicator = gr.Textbox(
                            label="Running Status",
                            value="⚪ Idle",
                            interactive=False
                        )
                        
                        gr.Markdown("### 🔍 Found Collisions")
                        collisions_table = gr.Dataframe(
                            headers=["Timestamp", "Address", "Private Key", "Worker ID", "Attempts"],
                            datatype=["str", "str", "str", "number", "number"],
                            interactive=False,
                            wrap=True
                        )
                
                # Event handlers
                start_btn.click(
                    fn=finder.start_search,
                    inputs=[csv_file, num_workers, batch_size],
                    outputs=[status_text, total_attempts_display, found_count_display, status_indicator, collisions_table]
                )
                
                stop_btn.click(
                    fn=finder.stop_search,
                    inputs=[],
                    outputs=[status_text, status_indicator]
                )
                
                # Timer for auto-refresh
                timer = gr.Timer(value=2, active=True)
                timer.tick(
                    fn=finder.get_update_data,
                    inputs=[],
                    outputs=[total_attempts_display, found_count_display, collisions_table]
                )
            
            # Test Mode Tab
            with gr.Tab("🧪 Test Mode"):
                gr.Markdown("### 🧪 Test with Specific Address")
                gr.Markdown("Search for a specific Ethereum address to verify the finder works correctly.")
                
                test_address = gr.Textbox(
                    label="Target Address",
                    placeholder="0x771f4c697b35677b107f9ddc9cea0c2976a9a23e",
                    value="0x771f4c697b35677b107f9ddc9cea0c2976a9a23e"
                )
                
                with gr.Row():
                    test_workers = gr.Slider(
                        minimum=1,
                        maximum=multiprocessing.cpu_count(),
                        value=2,
                        step=1,
                        label="Workers for Test"
                    )
                    test_batch = gr.Slider(
                        minimum=50,
                        maximum=500,
                        value=100,
                        step=50,
                        label="Batch Size"
                    )
                
                with gr.Row():
                    test_start_btn = gr.Button("🧪 Start Test", variant="primary")
                    test_stop_btn = gr.Button("⏹️ Stop Test", variant="stop")
                
                test_status = gr.Textbox(label="Test Status", interactive=False)
                test_attempts = gr.Number(label="Test Attempts", value=0, interactive=False)
                test_found = gr.Number(label="Found", value=0, interactive=False)
                test_status_indicator = gr.Textbox(label="Status", value="⚪ Idle", interactive=False)
                test_collisions = gr.Dataframe(
                    headers=["Timestamp", "Address", "Private Key", "Worker ID", "Attempts"],
                    visible=False
                )
                
                def start_test_wrapper(address, workers, batch):
                    result = finder.start_search(None, int(workers), int(batch), test_address=address)
                    # Return all 5 values matching the outputs
                    return result[0], result[1], result[2], result[3], result[4]
                
                test_start_btn.click(
                    fn=start_test_wrapper,
                    inputs=[test_address, test_workers, test_batch],
                    outputs=[test_status, test_attempts, test_found, test_status_indicator, test_collisions]
                )
                
                test_stop_btn.click(
                    fn=finder.stop_search,
                    inputs=[],
                    outputs=[test_status, test_status_indicator]
                )
                
                # Timer for test mode updates
                def update_test_stats():
                    if finder.is_running:
                        attempts, found = finder.get_stats()
                        return attempts, found
                    return 0, 0
                
                test_timer = gr.Timer(value=1, active=True)
                test_timer.tick(
                    fn=update_test_stats,
                    inputs=[],
                    outputs=[test_attempts, test_found]
                )
            
            # Information Tab
            with gr.Tab("ℹ️ Information"):
                gr.Markdown("""
                ## 📚 How It Works
                
                This tool uses multi-processing to generate Ethereum private keys and check if their corresponding 
                public addresses match any addresses in your uploaded CSV file.
                
                ### 🔧 Features:
                - **Multi-processing**: Utilizes all CPU cores for maximum performance
                - **High Entropy**: Uses multiple entropy sources for cryptographically secure key generation
                - **Real-time Monitoring**: View progress and found collisions in real-time
                - **Automatic Saving**: All found collisions are saved to files automatically
                
                ### ⚠️ Important Notes:
                - **Statistical Impossibility**: Finding a collision is mathematically nearly impossible
                - **Educational Purpose**: This tool demonstrates the security of Ethereum's address space
                - **Resource Usage**: High CPU usage is expected during operation
                
                ### 📊 Performance:
                - Modern CPUs can generate millions of keys per second
                - The Ethereum address space is 2^160 (approximately 1.46 × 10^48)
                - Even at billions of keys per second, finding a collision would take longer than the age of the universe
                
                ### 🔐 Security:
                - All operations are performed locally
                - No data is sent over the network
                - Generated keys are cryptographically secure
                
                ### 🖥️ Access Information:
                - **Local URL**: http://localhost:7860
                - **Network URL**: http://127.0.0.1:7860
                - The application will automatically open in your default browser
                """)
        
    return app

# ============================================================================
# MAIN ENTRY POINT
# ============================================================================

if __name__ == "__main__":
    multiprocessing.freeze_support()
    
    # Create and launch the Gradio interface
    app = create_interface()
    
    print("=" * 60)
    print("🔑 Ethereum Address Collision Finder")
    print("=" * 60)
    print("\n🚀 Starting web interface...")
    print("📱 Local URL: http://localhost:7860")
    print("🌐 Network URL: http://127.0.0.1:7860")
    print("\n💡 The application will open in your default browser automatically")
    print("📋 Press Ctrl+C to stop the server\n")
    
    # Launch with proper settings
    app.launch(
        server_name="127.0.0.1",  # Use localhost instead of 0.0.0.0
        server_port=7860,
        share=False,
        theme=gr.themes.Soft(),
        inbrowser=True,  # Automatically open in browser
        show_error=True
    )