#!/usr/bin/env python3
"""
RinCoin Stratum Proxy Server - FIXED VERSION
Fixed block hash calculation and validation issues
DEBUG: we get node logs with 'docker logs --tail=200 rincoin-node'
"""

import socket
import threading
import json
import time
import requests
import hashlib
import struct
from requests.auth import HTTPBasicAuth

class RinCoinStratumProxy:
    def __init__(self, stratum_host='0.0.0.0', stratum_port=3334, 
                 rpc_host='127.0.0.1', rpc_port=9556, 
                 rpc_user='rinrpc', rpc_password='745ce784d5d537fc06105a1b935b7657903cfc71a5fb3b90',
                 target_address='rin1qahvvv9d5f3443wtckeqavwp9950wacxfmwv20q',
                 submit_all_blocks=False):
        
        self.stratum_host = stratum_host
        self.stratum_port = stratum_port
        self.rpc_host = rpc_host
        self.rpc_port = rpc_port
        self.rpc_user = rpc_user
        self.rpc_password = rpc_password
        self.target_address = target_address
        self.submit_all_blocks = submit_all_blocks  # If True, submit even invalid blocks for testing
        
        self.clients = {}
        self.job_counter = 0
        self.current_job = None
        self.running = True
        self.extranonce1_counter = 0
        
        # Dynamic difficulty adjustment
        self.share_stats = {}  # Track shares per client
        self.last_difficulty_adjustment = time.time()
        self.target_share_interval = 120  # Target: 1 share every 2 minutes per miner (aligned with ~1min blocks)
        
        # Production monitoring
        self.stats = {
            'start_time': time.time(),
            'total_shares': 0,
            'accepted_shares': 0,
            'rejected_shares': 0,
            'blocks_found': 0,
            'total_hashrate': 0,
            'connections': 0
        }
        self.max_connections = 50  # Production limit
        
        # Logging setup
        self.log_file = "mining_log.txt"
        self.init_log_file()
        
        print(f"RinCoin Stratum Proxy Server")
        print(f"Stratum: {stratum_host}:{stratum_port}")
        print(f"RPC: {rpc_host}:{rpc_port}")
        print(f"Target: {target_address}")
        
    def init_log_file(self):
        """Initialize mining log file with header"""
        try:
            with open(self.log_file, 'w') as f:
                f.write("=" * 80 + "\n")
                f.write("RinCoin Mining Log\n")
                f.write("=" * 80 + "\n")
                f.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"Target Address: {self.target_address}\n")
                f.write(f"Stratum: {self.stratum_host}:{self.stratum_port}\n")
                f.write(f"RPC: {self.rpc_host}:{self.rpc_port}\n")
                f.write("=" * 80 + "\n\n")
        except Exception as e:
            print(f"Failed to initialize log file: {e}")
    
    def log_hash_found(self, hash_hex, difficulty, height, reward, nonce, ntime):
        """Log found hash to file"""
        try:
            timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
            with open(self.log_file, 'a') as f:
                f.write(f"[{timestamp}] 🎉 HASH FOUND!\n")
                f.write(f"   Hash: {hash_hex}\n")
                f.write(f"   Difficulty: {difficulty:.6f}\n")
                f.write(f"   Height: {height}\n")
                f.write(f"   Reward: {reward:.8f} RIN\n")
                f.write(f"   Nonce: {nonce}\n")
                f.write(f"   Time: {ntime}\n")
                f.write("-" * 40 + "\n")
        except Exception as e:
            print(f"Failed to log hash: {e}")
    
    def log_wallet_balance(self):
        """Log current wallet balance to file"""
        try:
            balance = self.rpc_call("getbalance")
            if balance is not None:
                timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
                with open(self.log_file, 'a') as f:
                    f.write(f"[{timestamp}] 💰 Wallet Balance: {balance:.8f} RIN\n")
        except Exception as e:
            print(f"Failed to log wallet balance: {e}")
    
    def rpc_call(self, method, params=[]):
        """Make RPC call to RinCoin node"""
        try:
            url = f"http://{self.rpc_host}:{self.rpc_port}/"
            headers = {'content-type': 'text/plain'}
            auth = HTTPBasicAuth(self.rpc_user, self.rpc_password)
            
            payload = {
                "jsonrpc": "1.0",
                "id": "stratum_proxy",
                "method": method,
                "params": params
            }
            
            response = requests.post(url, json=payload, headers=headers, auth=auth, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                if 'error' in result and result['error'] is not None:
                    print(f"RPC Error: {result['error']}")
                    return None
                return result.get('result')
            else:
                print(f"HTTP Error: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"RPC Call Error: {e}")
            return None

    def encode_varint(self, n):
        """Encode integer as Bitcoin-style varint"""
        if n < 0xfd:
            return bytes([n])
        elif n <= 0xffff:
            return b"\xfd" + struct.pack('<H', n)
        elif n <= 0xffffffff:
            return b"\xfe" + struct.pack('<I', n)
        else:
            return b"\xff" + struct.pack('<Q', n)

    def decode_bech32_address(self, address):
        """Decode RinCoin bech32 address to script"""
        try:
            if not address or not address.startswith('rin1'):
                raise ValueError("Not a RinCoin bech32 address")
            
            result = self.rpc_call("validateaddress", [address])
            if not result or not result.get('isvalid'):
                raise ValueError("Address not valid per node")
            script_hex = result.get('scriptPubKey')
            if not script_hex:
                raise ValueError("Node did not return scriptPubKey")
            return bytes.fromhex(script_hex)
        except Exception as e:
            print(f"Address decode error: {e}")
            return None

    def build_coinbase_transaction(self, template, extranonce1, extranonce2):
        """Build coinbase transaction variants (with and without witness) for default address"""
        return self.build_coinbase_transaction_for_address(template, extranonce1, extranonce2, self.target_address)
    
    def build_coinbase_transaction_for_address(self, template, extranonce1, extranonce2, target_address):
        """Build coinbase transaction variants (with and without witness)"""
        try:
            has_witness_commitment = template.get('default_witness_commitment') is not None

            # Common parts
            value = template.get('coinbasevalue', 0)
            script_pubkey = self.decode_bech32_address(target_address)
            if not script_pubkey:
                return None, None
            witness_commitment = template.get('default_witness_commitment')

            # ScriptSig (block height minimal push + tag + extranonces)
            height = template.get('height', 0)
            height_bytes = struct.pack('<I', height)
            height_compact = bytes([len(height_bytes.rstrip(b'\x00'))]) + height_bytes.rstrip(b'\x00')
            scriptsig = height_compact + b'/RinCoin/' + extranonce1.encode() + extranonce2.encode()

            # Helper to build outputs blob
            def build_outputs_blob() -> bytes:
                outputs_blob = b''
                outputs_list = []
                # Main output
                outputs_list.append(struct.pack('<Q', value) + self.encode_varint(len(script_pubkey)) + script_pubkey)
                # Witness commitment OP_RETURN output if present
                if witness_commitment:
                    commit_script = bytes.fromhex(witness_commitment)
                    outputs_list.append(struct.pack('<Q', 0) + self.encode_varint(len(commit_script)) + commit_script)
                outputs_blob += self.encode_varint(len(outputs_list))
                for out in outputs_list:
                    outputs_blob += out
                return outputs_blob

            # Build non-witness serialization (txid serialization)
            cb_nowit = b''
            cb_nowit += struct.pack('<I', 1)               # version
            cb_nowit += b'\x01'                            # input count
            cb_nowit += b'\x00' * 32                       # prevout hash
            cb_nowit += b'\xff\xff\xff\xff'             # prevout index
            cb_nowit += self.encode_varint(len(scriptsig)) + scriptsig
            cb_nowit += b'\xff\xff\xff\xff'             # sequence
            cb_nowit += build_outputs_blob()                # outputs
            cb_nowit += struct.pack('<I', 0)               # locktime

            # Build with-witness serialization (block serialization)
            if has_witness_commitment:
                cb_wit = b''
                cb_wit += struct.pack('<I', 1)             # version
                cb_wit += b'\x00\x01'                     # segwit marker+flag
                cb_wit += b'\x01'                          # input count
                cb_wit += b'\x00' * 32                     # prevout hash
                cb_wit += b'\xff\xff\xff\xff'           # prevout index
                cb_wit += self.encode_varint(len(scriptsig)) + scriptsig
                cb_wit += b'\xff\xff\xff\xff'           # sequence
                cb_wit += build_outputs_blob()              # outputs
                # Witness stack for coinbase (32-byte reserved value)
                cb_wit += b'\x01'                          # witness stack count
                cb_wit += b'\x20'                          # item length
                cb_wit += b'\x00' * 32                     # reserved value
                cb_wit += struct.pack('<I', 0)             # locktime
            else:
                cb_wit = cb_nowit

            return cb_wit, cb_nowit

        except Exception as e:
            print(f"Coinbase construction error: {e}")
            return None, None

    def calculate_merkle_root(self, coinbase_txid, transactions):
        """Calculate merkle root with coinbase at index 0"""
        try:
            # Start with all transaction hashes (coinbase + others)
            hashes = [coinbase_txid]
            for tx in transactions:
                hashes.append(bytes.fromhex(tx['hash'])[::-1])  # Reverse for little-endian

            # Build merkle tree
            while len(hashes) > 1:
                if len(hashes) % 2 == 1:
                    hashes.append(hashes[-1])  # Duplicate last hash if odd

                next_level = []
                for i in range(0, len(hashes), 2):
                    combined = hashes[i] + hashes[i + 1]
                    next_level.append(hashlib.sha256(hashlib.sha256(combined).digest()).digest())

                hashes = next_level

            return hashes[0] if hashes else b'\x00' * 32

        except Exception as e:
            print(f"Merkle root calculation error: {e}")
            return b'\x00' * 32

    def calculate_merkle_branches(self, tx_hashes, tx_index):
        """Calculate merkle branches for a specific transaction index"""
        try:
            if not tx_hashes or tx_index >= len(tx_hashes):
                return []

            branches = []
            current_index = tx_index

            # Start with the full list of transaction hashes
            hashes = tx_hashes[:]

            while len(hashes) > 1:
                if len(hashes) % 2 == 1:
                    hashes.append(hashes[-1])  # Duplicate last hash if odd

                # Find the partner for current_index
                partner_index = current_index ^ 1  # Flip the least significant bit

                if partner_index < len(hashes):
                    # Add the partner hash to branches (in big-endian for stratum)
                    branches.append(hashes[partner_index][::-1].hex())
                else:
                    # This shouldn't happen, but add the duplicate if it does
                    branches.append(hashes[current_index][::-1].hex())

                # Move to next level
                next_level = []
                for i in range(0, len(hashes), 2):
                    combined = hashes[i] + hashes[i + 1]
                    next_level.append(hashlib.sha256(hashlib.sha256(combined).digest()).digest())

                hashes = next_level
                current_index //= 2

            return branches

        except Exception as e:
            print(f"Merkle branches calculation error: {e}")
            return []

    def bits_to_target(self, bits_hex):
        """Convert bits to target - FIXED VERSION"""
        try:
            bits = int(bits_hex, 16)
            exponent = bits >> 24
            mantissa = bits & 0xffffff
            
            # Bitcoin target calculation
            if exponent <= 3:
                target = mantissa >> (8 * (3 - exponent))
            else:
                target = mantissa << (8 * (exponent - 3))
            
            return f"{target:064x}"
        except Exception as e:
            print(f"Bits to target error: {e}")
            return "0000ffff00000000000000000000000000000000000000000000000000000000"

    def get_block_template(self):
        """Get new block template and create Stratum job"""
        try:
            template = self.rpc_call("getblocktemplate", [{"rules": ["segwit", "mweb"]}])
            if not template:
                return None
                
            self.job_counter += 1
            
            job = {
                "job_id": f"job_{self.job_counter:08x}",
                "template": template,
                "prevhash": template.get("previousblockhash", "0" * 64),
                "version": template.get('version', 1),
                "bits": template.get('bits', '1d00ffff'),
                "ntime": f"{int(time.time()):08x}",
                "target": self.bits_to_target(template.get('bits', '1d00ffff')),
                "height": template.get('height', 0),
                "coinbasevalue": template.get('coinbasevalue', 0),
                "transactions": template.get('transactions', [])
            }
            
            self.current_job = job
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            network_difficulty = self.calculate_network_difficulty(job['target'])
            print(f"[{timestamp}] 🆕 NEW JOB: {job['job_id']} | Height: {job['height']} | Reward: {job['coinbasevalue']/100000000:.2f} RIN")
            print(f"   🎯 Network Difficulty: {network_difficulty:.6f} | Bits: {job['bits']}")
            print(f"   📍 Target: {job['target'][:16]}... | Transactions: {len(job['transactions'])}")
            return job
            
        except Exception as e:
            print(f"Get block template error: {e}")
            return None

    def calculate_share_difficulty(self, hash_hex, target_hex):
        """Calculate actual share difficulty from hash - FIXED"""
        try:
            hash_int = int(hash_hex, 16)
            
            if hash_int == 0:
                return float('inf')  # Perfect hash
                
            # Bitcoin-style difficulty calculation using difficulty 1 target
            # Difficulty 1 target for mainnet
            diff1_target = 0x00000000FFFF0000000000000000000000000000000000000000000000000000
            
            # Share difficulty = how much harder this hash was compared to diff 1
            difficulty = diff1_target / hash_int
            
            return difficulty
        except Exception as e:
            print(f"Difficulty calculation error: {e}")
            return 0.0

    def calculate_network_difficulty(self, target_hex):
        """Calculate network difficulty from target - FIXED"""
        try:
            target_int = int(target_hex, 16)
            
            # Bitcoin difficulty 1.0 target
            diff1_target = 0x00000000FFFF0000000000000000000000000000000000000000000000000000
            
            # Network difficulty = how much harder than difficulty 1.0
            network_difficulty = diff1_target / target_int
            
            return network_difficulty
        except Exception as e:
            print(f"Network difficulty calculation error: {e}")
            return 1.0

    def calculate_optimal_difficulty(self, addr, is_new_client=False):
        """Calculate optimal stratum difficulty for a client based on hashrate and network conditions"""
        try:
            # Get current network difficulty
            network_diff = self.calculate_network_difficulty(self.current_job['target']) if self.current_job else 1.0
            
            if is_new_client:
                # For new clients, start with a conservative estimate
                # Assuming typical CPU miner does 500kH/s, aim for 1 share per 30 seconds
                # Difficulty calculation: target_shares_per_second = hashrate / (difficulty * 2^32 / 65535)
                # Rearranged: difficulty = hashrate / (target_shares_per_second * 2^32 / 65535)
                estimated_hashrate = 500000  # 500 kH/s conservative estimate
                target_shares_per_second = 1.0 / self.target_share_interval
                
                # Bitcoin-style difficulty calculation
                diff1_target = 0x00000000FFFF0000000000000000000000000000000000000000000000000000
                optimal_difficulty = (estimated_hashrate * diff1_target) / (target_shares_per_second * (2**256))
                
                # Calculate proper difficulty based on expected hashrate for solo mining
                # Target: 1 share every 2 minutes (120 seconds)
                # Conservative estimate: 500 kH/s for CPU mining
                estimated_hashrate = 500000  # 500 kH/s
                target_shares_per_second = 1.0 / self.target_share_interval
                
                # Calculate difficulty to achieve target share rate
                # Formula: difficulty = hashrate / (target_shares_per_second * 2^32 / max_target)
                max_target = 0x00000000FFFF0000000000000000000000000000000000000000000000000000
                optimal_difficulty = (estimated_hashrate * target_shares_per_second) / (max_target / (2**32))
                
                # Ensure it's not too low or too high
                optimal_difficulty = max(0.0000001, min(optimal_difficulty, network_diff * 0.01))
                
                print(f"   📊 New client {addr}: Network diff {network_diff:.6f}, starting with {optimal_difficulty:.6f}")
                return optimal_difficulty
            
            # For existing clients, adjust based on actual performance
            client_data = self.clients.get(addr, {})
            if not client_data:
                return self.calculate_optimal_difficulty(addr, is_new_client=True)
            
            current_time = time.time()
            connect_time = client_data.get('connect_time', current_time)
            share_count = client_data.get('share_count', 0)
            current_difficulty = client_data.get('stratum_difficulty', 0.001)
            
            # Calculate actual share rate
            mining_duration = current_time - connect_time
            if mining_duration < 60:  # Need at least 1 minute of data
                return current_difficulty
            
            actual_share_rate = share_count / mining_duration  # shares per second
            target_share_rate = 1.0 / self.target_share_interval
            
            # Adjust difficulty to hit target share rate
            if actual_share_rate > 0:
                difficulty_multiplier = actual_share_rate / target_share_rate
                new_difficulty = current_difficulty * difficulty_multiplier
                
                # Apply conservative adjustment (don't change too dramatically)
                max_change = 2.0  # Don't change by more than 2x in one adjustment
                if difficulty_multiplier > max_change:
                    new_difficulty = current_difficulty * max_change
                elif difficulty_multiplier < (1.0 / max_change):
                    new_difficulty = current_difficulty / max_change
                
                # Bounds checking for solo mining  
                min_difficulty = 0.00000001  # Extremely low minimum for CPU mining
                max_difficulty = 0.0001       # Maximum 0.01% of typical network diff
                new_difficulty = max(min_difficulty, min(max_difficulty, new_difficulty))
                
                print(f"   📊 {addr}: {share_count} shares in {mining_duration:.0f}s ({actual_share_rate:.4f}/s), adjusting {current_difficulty:.6f} → {new_difficulty:.6f}")
                return new_difficulty
            
            return current_difficulty
            
        except Exception as e:
            print(f"Difficulty calculation error: {e}")
            # Fallback to conservative difficulty
            network_diff = self.calculate_network_difficulty(self.current_job['target']) if self.current_job else 1.0
            return max(network_diff * 0.001, 0.0001)

    def adjust_client_difficulty(self, addr):
        """Adjust difficulty for a specific client if needed"""
        try:
            current_time = time.time()
            
            # Only adjust every 2 minutes minimum
            last_adjustment = self.clients[addr].get('last_difficulty_adjustment', 0)
            if current_time - last_adjustment < 120:
                return
            
            new_difficulty = self.calculate_optimal_difficulty(addr, is_new_client=False)
            current_difficulty = self.clients[addr].get('stratum_difficulty', 0.001)
            
            # Only send update if difficulty changed significantly (>20%)
            if abs(new_difficulty - current_difficulty) / current_difficulty > 0.2:
                self.clients[addr]['stratum_difficulty'] = new_difficulty
                self.clients[addr]['last_difficulty_adjustment'] = current_time
                
                # Send new difficulty to client
                if 'socket' in self.clients[addr]:
                    self.send_stratum_notification(self.clients[addr]['socket'], "mining.set_difficulty", [new_difficulty])
                    print(f"   🎯 Updated difficulty for {addr}: {current_difficulty:.6f} → {new_difficulty:.6f}")
                
        except Exception as e:
            print(f"Difficulty adjustment error for {addr}: {e}")

    def submit_share(self, job, extranonce1, extranonce2, ntime, nonce, addr=None, target_address=None):
        """Validate share and submit block if valid - FIXED VERSION"""
        try:
            # Use provided address or default
            address = target_address or self.target_address
            
            # Build coinbase (with and without witness)
            coinbase_wit, coinbase_nowit = self.build_coinbase_transaction_for_address(
                job['template'], extranonce1, extranonce2, address)
            if not coinbase_wit or not coinbase_nowit:
                return False, "Coinbase construction failed"
            
            # Calculate coinbase txid (non-witness serialization)
            coinbase_txid = hashlib.sha256(hashlib.sha256(coinbase_nowit).digest()).digest()[::-1]
            
            # Calculate merkle root
            merkle_root = self.calculate_merkle_root(coinbase_txid, job['transactions'])
            
            # Build block header - FIXED ENDIANNESS
            header = b''
            header += struct.pack('<I', job['version'])                    # Version (little-endian)
            header += bytes.fromhex(job['prevhash'])[::-1]                # Previous block hash (big-endian in block)
            header += merkle_root                                         # Merkle root (already in correct endian)
            header += struct.pack('<I', int(ntime, 16))                  # Timestamp (little-endian)
            header += bytes.fromhex(job['bits'])                         # Bits (big-endian in block)
            header += struct.pack('<I', int(nonce, 16))                  # Nonce (little-endian)
            
            # Calculate block hash - FIXED DOUBLE SHA256
            block_hash = hashlib.sha256(hashlib.sha256(header).digest()).digest()
            block_hash_hex = block_hash[::-1].hex()  # Reverse for display/comparison
            
            # Calculate real difficulties
            share_difficulty = self.calculate_share_difficulty(block_hash_hex, job['target'])
            network_difficulty = self.calculate_network_difficulty(job['target'])
            
            # Check if hash meets stratum difficulty first, then network target
            hash_int = int(block_hash_hex, 16)
            network_target_int = int(job['target'], 16)

            # Get the stratum difficulty that was sent to this miner
            client_stratum_diff = self.clients.get(addr, {}).get('stratum_difficulty', 0.001)
            diff1_target = 0x00000000FFFF0000000000000000000000000000000000000000000000000000
            # Correct stratum target calculation: target = diff1_target / difficulty
            stratum_target_int = int(diff1_target / client_stratum_diff)

            meets_stratum_target = hash_int <= stratum_target_int
            meets_network_target = hash_int <= network_target_int
            
            # CRITICAL FIX: Check if share difficulty meets minimum requirements
            meets_stratum_difficulty = share_difficulty >= client_stratum_diff
            meets_network_difficulty = share_difficulty >= network_difficulty

            # Enhanced logging
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            network_percentage = (share_difficulty / network_difficulty) * 100 if network_difficulty > 0 else 0
            stratum_percentage = (share_difficulty / client_stratum_diff) * 100 if client_stratum_diff > 0 else 0

            # Progress indicator based on difficulty comparison
            if meets_network_difficulty:
                progress_icon = "🎉"  # Block found!
            elif meets_stratum_difficulty:
                progress_icon = "✅"  # Valid share for stratum
            elif network_percentage >= 50:
                progress_icon = "🔥"  # Very close to network
            elif network_percentage >= 10:
                progress_icon = "⚡"  # Getting warm
            elif network_percentage >= 1:
                progress_icon = "💫"  # Some progress
            else:
                progress_icon = "📊"  # Low progress

            print(f"[{timestamp}] {progress_icon} SHARE: job={job['job_id']} | nonce={nonce} | hash={block_hash_hex[:16]}...")
            print(f"   🎯 Share Diff: {share_difficulty:.2e} | Stratum Diff: {client_stratum_diff:.6f} | Network Diff: {network_difficulty:.6f}")
            print(f"   📈 Progress: {network_percentage:.4f}% of network, {stratum_percentage:.1f}% of stratum")
            print(f"   📍 Target: {job['target'][:16]}... | Height: {job['height']}")
            print(f"   ⏰ Time: {ntime} | Extranonce: {extranonce1}:{extranonce2}")
            print(f"   🔍 Difficulty Check: Share {share_difficulty:.2e} vs Stratum {client_stratum_diff:.6f} = {'✅' if meets_stratum_difficulty else '❌'}")
            print(f"   🔍 Difficulty Check: Share {share_difficulty:.2e} vs Network {network_difficulty:.6f} = {'✅' if meets_network_difficulty else '❌'}")

            # Handle submit_all_blocks mode - bypass all difficulty checks
            if self.submit_all_blocks:
                print(f"   🧪 TEST MODE: Submitting ALL shares to node for validation")
                # Update stats for test mode
                self.stats['total_shares'] += 1
                self.stats['accepted_shares'] += 1
                
                # Always submit in test mode
                should_submit = True
            else:
                # Normal mode - check stratum difficulty properly
                if not meets_stratum_difficulty:
                    # Share doesn't meet minimum stratum difficulty - reject
                    print(f"   ❌ Share rejected (difficulty too low: {share_difficulty:.2e} < {client_stratum_diff:.6f})")
                    self.stats['total_shares'] += 1
                    self.stats['rejected_shares'] += 1
                    return False, "Share does not meet stratum difficulty"

                # Valid stratum share! Update client stats
                if addr and addr in self.clients:
                    self.clients[addr]['share_count'] = self.clients[addr].get('share_count', 0) + 1
                    self.clients[addr]['last_share_time'] = time.time()
                    
                    # Update global stats
                    self.stats['total_shares'] += 1
                    self.stats['accepted_shares'] += 1
                    
                    # Check if we should adjust difficulty
                    self.adjust_client_difficulty(addr)
                
                # Only submit if it meets network difficulty in normal mode
                should_submit = meets_network_difficulty

            # Submit block if conditions are met
            if should_submit:
                if meets_network_difficulty:
                    print(f"   🎉 VALID BLOCK FOUND! Hash: {block_hash_hex}")
                else:
                    print(f"   🧪 TEST BLOCK SUBMISSION! Hash: {block_hash_hex} (submit_all_blocks=True)")
                print(f"   💰 Reward: {job['coinbasevalue']/100000000:.2f} RIN -> {address}")
                print(f"   📊 Block height: {job['height']}")
                print(f"   🔍 Difficulty: {share_difficulty:.6f} (target: {network_difficulty:.6f})")

                # Log the found hash
                reward_rin = job['coinbasevalue'] / 100000000
                self.log_hash_found(block_hash_hex, share_difficulty, job['height'], reward_rin, nonce, ntime)

                # Build complete block
                block = header

                # Transaction count
                tx_count = 1 + len(job['transactions'])
                block += self.encode_varint(tx_count)

                # Add coinbase transaction (witness variant for block body)
                block += coinbase_wit

                # Add other transactions
                for tx in job['transactions']:
                    block += bytes.fromhex(tx['data'])

                # Submit block
                block_hex = block.hex()
                print(f"   📦 Submitting block of size {len(block_hex)//2} bytes...")

                result = self.rpc_call("submitblock", [block_hex])

                if result is None:
                    print(f"   ✅ Block accepted by network!")
                    # Update block stats
                    self.stats['blocks_found'] += 1
                    # Log wallet balance after successful block submission
                    self.log_wallet_balance()
                    return True, "Block found and submitted"
                else:
                    print(f"   ❌ Block rejected: {result}")
                    print(f"   🔍 Debug: Block size {len(block_hex)//2} bytes, {len(job['transactions'])} transactions")
                    return False, f"Block rejected: {result}"
            else:
                # Valid stratum share but not network-valid block (normal mode only)
                if not self.submit_all_blocks:
                    print(f"   ✅ Valid stratum share accepted")
                    return True, "Valid stratum share"
                else:
                    # This shouldn't happen in test mode since should_submit would be True
                    return True, "Test mode share processed"
                
        except Exception as e:
            print(f"Share submission error: {e}")
            return False, f"Submission error: {e}"

    def send_stratum_response(self, client, msg_id, result, error=None):
        """Send Stratum response to client"""
        try:
            response = {
                "id": msg_id,
                "result": result,
                "error": error
            }
            message = json.dumps(response) + "\n"
            client.send(message.encode('utf-8'))
        except Exception as e:
            print(f"Send response error: {e}")

    def send_stratum_notification(self, client, method, params):
        """Send Stratum notification to client"""
        try:
            notification = {
                "id": None,
                "method": method,
                "params": params
            }
            message = json.dumps(notification) + "\n"
            client.send(message.encode('utf-8'))
        except Exception as e:
            print(f"Send notification error: {e}")

    def handle_stratum_message(self, client, addr, message):
        """Handle incoming Stratum message from miner"""
        try:
            data = json.loads(message.strip())
            method = data.get("method")
            msg_id = data.get("id")
            params = data.get("params", [])
            
            if method == "mining.subscribe":
                # Generate unique extranonce1 for this connection
                self.extranonce1_counter += 1
                extranonce1 = f"{self.extranonce1_counter:08x}"
                
                # Store extranonce1 for this client
                if addr not in self.clients:
                    self.clients[addr] = {}
                self.clients[addr]['extranonce1'] = extranonce1
                
                # Subscribe response
                self.send_stratum_response(client, msg_id, [
                    [["mining.set_difficulty", "subscription_id"], ["mining.notify", "subscription_id"]],
                    extranonce1,
                    4  # extranonce2 size
                ])
                
                # Calculate optimal stratum difficulty for this client
                initial_difficulty = self.calculate_optimal_difficulty(addr, is_new_client=True)
                # Store stratum difficulty for this client
                if addr not in self.clients:
                    self.clients[addr] = {}
                self.clients[addr]['stratum_difficulty'] = initial_difficulty
                self.clients[addr]['last_share_time'] = time.time()
                self.clients[addr]['share_count'] = 0
                self.clients[addr]['connect_time'] = time.time()
                
                self.send_stratum_notification(client, "mining.set_difficulty", [initial_difficulty])
                print(f"   🎯 Set initial difficulty {initial_difficulty:.6f} for {addr}")
                
                # Send initial job
                if self.current_job:
                    self.send_job_to_client(client, self.current_job)
                else:
                    if self.get_block_template():
                        self.send_job_to_client(client, self.current_job)
            
            elif method == "mining.authorize":
                username = params[0] if params else "anonymous"
                self.clients[addr]['username'] = username
                self.send_stratum_response(client, msg_id, True)
                timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
                print(f"[{timestamp}] 🔐 [{addr}] Authorized as {username}")
            
            elif method == "mining.extranonce.subscribe":
                self.send_stratum_response(client, msg_id, True)
            
            elif method == "mining.submit":
                if len(params) >= 5:
                    username = params[0]
                    job_id = params[1]
                    extranonce2 = params[2]
                    ntime = params[3]
                    nonce = params[4]
                    
                    print(f"[{addr}] Submit: {username} | job={job_id} | nonce={nonce}")
                    
                    # Always validate against current job
                    if self.current_job:
                        extranonce1 = self.clients[addr].get('extranonce1', '00000000')
                        
                        # Submit share
                        success, message = self.submit_share(self.current_job, extranonce1, extranonce2, ntime, nonce, addr)
                        
                        # Always accept shares for debugging, even if they don't meet target
                        self.send_stratum_response(client, msg_id, True)
                        
                        if success and "Block found" in message:
                            # Get new job after block found
                            threading.Thread(target=self.update_job_after_block, daemon=True).start()
                    else:
                        self.send_stratum_response(client, msg_id, True)
                else:
                    self.send_stratum_response(client, msg_id, False, "Invalid parameters")
            
            else:
                print(f"[{addr}] Unknown method: {method}")
                self.send_stratum_response(client, msg_id, None, "Unknown method")
                
        except json.JSONDecodeError:
            print(f"[{addr}] Invalid JSON: {message}")
        except Exception as e:
            print(f"[{addr}] Message handling error: {e}")

    def send_job_to_client(self, client, job):
        """Send mining job to specific client with proper stratum parameters"""
        try:
            # Get network difficulty for display
            network_difficulty = self.calculate_network_difficulty(job['target'])
            
            # Build proper coinbase transaction parts for stratum protocol
            template = job.get('template', {})
            height = job.get('height', 0)
            
            # Encode height as minimal push (BIP34 compliance)
            if height == 0:
                height_push = b''
            else:
                height_bytes = struct.pack('<I', height)
                # Remove trailing zero bytes
                while len(height_bytes) > 1 and height_bytes[-1] == 0:
                    height_bytes = height_bytes[:-1]
                height_push = bytes([len(height_bytes)]) + height_bytes
            
            # Build coinbase input script: height + arbitrary data + extranonces
            coinb1_script = height_push + b'/RinCoin/'
            
            # Build coinbase transaction structure
            # Version (4 bytes)
            coinb1 = struct.pack('<I', 1)
            # Input count (1 byte) 
            coinb1 += b'\x01'
            # Previous output hash (32 bytes of zeros)
            coinb1 += b'\x00' * 32
            # Previous output index (4 bytes of 0xFF)
            coinb1 += b'\xff\xff\xff\xff'
            # Script length (will include extranonces)
            script_len_pos = len(coinb1)
            coinb1 += b'\x00'  # Placeholder for script length
            # Script data up to extranonce1
            coinb1 += coinb1_script
            
            # coinb2: everything after extranonces
            # For now, just the sequence and outputs
            coinb2 = b''
            coinb2 += b'\xff\xff\xff\xff'  # Sequence
            
            # Outputs
            script_pubkey = self.decode_bech32_address(self.target_address)
            if script_pubkey:
                value = template.get('coinbasevalue', 0)
                # Output count
                output_count = 1
                witness_commitment = template.get('default_witness_commitment')
                if witness_commitment:
                    output_count = 2
                
                coinb2 += self.encode_varint(output_count)
                
                # Main output (to mining address)
                coinb2 += struct.pack('<Q', value)
                coinb2 += self.encode_varint(len(script_pubkey))
                coinb2 += script_pubkey
                
                # Witness commitment output if present
                if witness_commitment:
                    commit_script = bytes.fromhex(witness_commitment)
                    coinb2 += struct.pack('<Q', 0)  # Zero value
                    coinb2 += self.encode_varint(len(commit_script))
                    coinb2 += commit_script
            
            # Locktime
            coinb2 += struct.pack('<I', 0)
            
            # Calculate total script length (coinb1_script + extranonces lengths)
            # extranonce1 is 8 hex chars (4 bytes), extranonce2 is 8 hex chars (4 bytes)
            total_script_len = len(coinb1_script) + 4 + 4  # extranonce1 + extranonce2 (binary)
            
            # Update script length in coinb1
            coinb1_list = list(coinb1)
            if total_script_len < 253:
                coinb1_list[script_len_pos] = total_script_len
            else:
                # Handle larger script lengths with varint encoding
                coinb1 = coinb1[:script_len_pos] + self.encode_varint(total_script_len) + coinb1[script_len_pos+1:]
            coinb1 = bytes(coinb1_list) if total_script_len < 253 else coinb1
            
            # Calculate merkle branches
            tx_hashes = []
            for tx in job.get('transactions', []):
                tx_hashes.append(bytes.fromhex(tx['hash'])[::-1])  # Little-endian
            
            merkle_branches = []
            if tx_hashes:
                # For coinbase at index 0, calculate merkle path
                merkle_branches = self.calculate_merkle_branches([b'\x00' * 32] + tx_hashes, 0)

            self.send_stratum_notification(client, "mining.notify", [
                job["job_id"],
                job["prevhash"],
                coinb1.hex(),
                coinb2.hex(),
                merkle_branches,
                f"{job['version']:08x}",
                job["bits"],
                job["ntime"],
                True  # clean_jobs
            ])
            
            print(f"   📤 Sent job to miner: Height {height} | NetDiff {network_difficulty:.6f} | Txs {len(job.get('transactions', []))}")
            
        except Exception as e:
            print(f"Failed to send job: {e}")
            import traceback
            traceback.print_exc()

    def update_job_after_block(self):
        """Update job after a block is found"""
        time.sleep(2)  # Brief delay to let network propagate
        if self.get_block_template():
            self.broadcast_new_job()

    def broadcast_new_job(self):
        """Broadcast new job to all connected clients"""
        if not self.current_job:
            return
            
        print(f"Broadcasting new job to {len(self.clients)} clients")
        
        for addr, client_data in list(self.clients.items()):
            try:
                if 'socket' in client_data:
                    self.send_job_to_client(client_data['socket'], self.current_job)
            except Exception as e:
                print(f"Failed to send job to {addr}: {e}")

    def handle_client(self, client, addr):
        """Handle individual client connection"""
        # Check connection limits
        if len(self.clients) >= self.max_connections:
            print(f"[{addr}] Connection rejected - max connections ({self.max_connections}) reached")
            client.close()
            return
            
        print(f"[{addr}] Connected (clients: {len(self.clients) + 1}/{self.max_connections})")
        if addr not in self.clients:
            self.clients[addr] = {}
        self.clients[addr]['socket'] = client
        self.stats['connections'] = len(self.clients)
        
        try:
            while self.running:
                data = client.recv(4096)
                if not data:
                    break
                
                # Handle multiple messages in one packet
                messages = data.decode('utf-8').strip().split('\n')
                for message in messages:
                    if message:
                        self.handle_stratum_message(client, addr, message)
                        
        except Exception as e:
            print(f"[{addr}] Client error: {e}")
        finally:
            client.close()
            if addr in self.clients:
                del self.clients[addr]
            self.stats['connections'] = len(self.clients)
            print(f"[{addr}] Disconnected (clients: {len(self.clients)}/{self.max_connections})")

    def print_stats(self):
        """Print pool statistics"""
        try:
            uptime = time.time() - self.stats['start_time']
            uptime_str = f"{int(uptime//3600)}h {int((uptime%3600)//60)}m"
            
            accept_rate = 0
            if self.stats['total_shares'] > 0:
                accept_rate = (self.stats['accepted_shares'] / self.stats['total_shares']) * 100
            
            print(f"\n📊 POOL STATISTICS:")
            print(f"   ⏰ Uptime: {uptime_str}")
            print(f"   👥 Connected miners: {self.stats['connections']}")
            print(f"   📈 Shares: {self.stats['accepted_shares']}/{self.stats['total_shares']} ({accept_rate:.1f}% accepted)")
            print(f"   🎉 Blocks found: {self.stats['blocks_found']}")
            print(f"   🎯 Network difficulty: {self.calculate_network_difficulty(self.current_job['target']) if self.current_job else 'unknown':.6f}")
            
            # Calculate hashrate estimate from connected clients
            total_hashrate = 0
            for addr, client_data in self.clients.items():
                share_count = client_data.get('share_count', 0)
                connect_time = client_data.get('connect_time', time.time())
                mining_duration = time.time() - connect_time
                if mining_duration > 60 and share_count > 0:
                    stratum_diff = client_data.get('stratum_difficulty', 0.001)
                    # Rough hashrate estimate: shares * difficulty * 2^32 / time
                    hashrate = (share_count * stratum_diff * 4294967296) / mining_duration
                    total_hashrate += hashrate
                    print(f"   🔥 {addr}: ~{hashrate/1000:.0f} kH/s")
            
            if total_hashrate > 0:
                print(f"   🚀 Total pool hashrate: ~{total_hashrate/1000:.0f} kH/s")
            print()
            
        except Exception as e:
            print(f"Stats error: {e}")

    def job_updater(self):
        """Periodically update mining jobs"""
        balance_log_counter = 0
        stats_counter = 0
        
        while self.running:
            try:
                time.sleep(30)  # Update every 30 seconds
                
                old_height = self.current_job['height'] if self.current_job else 0
                
                if self.get_block_template():
                    new_height = self.current_job['height']
                    if new_height > old_height:
                        print(f"New block detected! Broadcasting new job...")
                        self.broadcast_new_job()
                
                # Log wallet balance every 10 minutes (20 cycles of 30 seconds)
                balance_log_counter += 1
                if balance_log_counter >= 20:
                    self.log_wallet_balance()
                    balance_log_counter = 0
                
                # Print stats every 5 minutes (10 cycles of 30 seconds)
                stats_counter += 1
                if stats_counter >= 10:
                    self.print_stats()
                    stats_counter = 0
                        
            except Exception as e:
                print(f"Job updater error: {e}")

    def validate_mining_capability(self):
        """Validate that we can mine valid blocks against the RinCoin node"""
        try:
            print("🔍 Validating mining capability...")
            
            # Get block template
            template = self.rpc_call("getblocktemplate", [{"rules": ["segwit", "mweb"]}])
            if not template:
                print("❌ Cannot get block template")
                return False
            
            # Test address validation
            result = self.rpc_call("validateaddress", [self.target_address])
            if not result or not result.get('isvalid'):
                print(f"❌ Target address {self.target_address} is invalid")
                return False
            
            print(f"✅ Target address {self.target_address} is valid")
            
            # Test coinbase construction
            coinbase_wit, coinbase_nowit = self.build_coinbase_transaction_for_address(
                template, "00000001", "01000000", self.target_address)
            
            if not coinbase_wit or not coinbase_nowit:
                print("❌ Cannot construct coinbase transaction")
                return False
            
            print(f"✅ Coinbase construction works")
            
            # Test block header construction with dummy values
            test_nonce = "12345678"
            test_ntime = f"{int(time.time()):08x}"
            
            # Calculate merkle root
            coinbase_txid = hashlib.sha256(hashlib.sha256(coinbase_nowit).digest()).digest()[::-1]
            merkle_root = self.calculate_merkle_root(coinbase_txid, template.get('transactions', []))
            
            # Build test header
            header = b''
            header += struct.pack('<I', template.get('version', 1))
            header += bytes.fromhex(template.get("previousblockhash", "0" * 64))[::-1]
            header += merkle_root
            header += struct.pack('<I', int(test_ntime, 16))
            header += bytes.fromhex(template.get('bits', '1d00ffff'))
            header += struct.pack('<I', int(test_nonce, 16))
            
            # Calculate hash
            block_hash = hashlib.sha256(hashlib.sha256(header).digest()).digest()
            block_hash_hex = block_hash[::-1].hex()
            
            print(f"✅ Block header construction works")
            print(f"   Test hash: {block_hash_hex}")
            
            # Check difficulty calculation
            target = self.bits_to_target(template.get('bits', '1d00ffff'))
            network_diff = self.calculate_network_difficulty(target)
            
            print(f"✅ Difficulty calculation works")
            print(f"   Network difficulty: {network_diff:.6f}")
            print(f"   Bits: {template.get('bits', '1d00ffff')}")
            print(f"   Target: {target[:16]}...")
            
            # Test RPC submission (dry run - don't actually submit)
            print(f"✅ Ready to mine valid blocks!")
            print(f"   Block height: {template.get('height', 0)}")
            print(f"   Reward: {template.get('coinbasevalue', 0) / 100000000:.2f} RIN")
            print(f"   Transactions: {len(template.get('transactions', []))}")
            
            return True
            
        except Exception as e:
            print(f"❌ Mining validation failed: {e}")
            import traceback
            traceback.print_exc()
            return False

    def start(self):
        """Start the Stratum proxy server"""
        try:
            # Test RPC connection
            blockchain_info = self.rpc_call("getblockchaininfo")
            if not blockchain_info:
                print("❌ Failed to connect to RinCoin node!")
                return
            
            print(f"✅ Connected to RinCoin node (block {blockchain_info.get('blocks', 'unknown')})")
            print(f"   Chain: {blockchain_info.get('chain', 'unknown')}")
            print(f"   Difficulty: {blockchain_info.get('difficulty', 'unknown')}")
            
            # Validate mining capability
            if not self.validate_mining_capability():
                print("❌ Mining validation failed - cannot continue")
                return
            
            # Log initial wallet balance
            self.log_wallet_balance()
            
            # Get initial block template
            if not self.get_block_template():
                print("❌ Failed to get initial block template!")
                return
            
            # Start job updater thread
            job_thread = threading.Thread(target=self.job_updater, daemon=True)
            job_thread.start()
            
            # Start Stratum server
            server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            server_socket.bind((self.stratum_host, self.stratum_port))
            server_socket.listen(10)
            
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            print(f"[{timestamp}] 🚀 Mining Stratum proxy ready!")
            print(f"   📡 Listening on {self.stratum_host}:{self.stratum_port}")
            print(f"   💰 Mining to: {self.target_address}")
            print(f"   📊 Current job: {self.current_job['job_id'] if self.current_job else 'None'}")
            print(f"   📝 Mining log: {self.log_file}")
            print("")
            print("   🔧 Miner command:")
            print(f"   ./cpuminer -a rinhash -o stratum+tcp://{self.stratum_host}:{self.stratum_port} -u worker1 -p x -t 4")
            print("")
            
            while self.running:
                try:
                    client, addr = server_socket.accept()
                    client_thread = threading.Thread(
                        target=self.handle_client, 
                        args=(client, addr), 
                        daemon=True
                    )
                    client_thread.start()
                except KeyboardInterrupt:
                    print("\nShutting down...")
                    self.running = False
                    break
                except Exception as e:
                    print(f"Server error: {e}")
                    
        except Exception as e:
            print(f"Failed to start server: {e}")
        finally:
            print("Server stopped")

if __name__ == "__main__":
    import sys
    
    # Check for submit_all_blocks parameter
    submit_all = False
    if len(sys.argv) > 1 and sys.argv[1] == "--submit-all-blocks":
        submit_all = True
        print("🧪 TEST MODE: Will submit ALL blocks for validation (submit_all_blocks=True)")
    
    proxy = RinCoinStratumProxy(submit_all_blocks=submit_all)
    proxy.start()