#!/usr/bin/env python3
"""
RinCoin Stratum Proxy Server - PRODUCTION VERSION
Bridges cpuminer-opt-rin (Stratum protocol) to RinCoin node (RPC protocol)
For real solo mining with actual block construction and submission
"""

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

class RinCoinStratumProxy:
    def __init__(self, stratum_host='0.0.0.0', stratum_port=3333, 
                 rpc_host='127.0.0.1', rpc_port=9556, 
                 rpc_user='rinrpc', rpc_password='745ce784d5d537fc06105a1b935b7657903cfc71a5fb3b90',
                 target_address='rin1qahvvv9d5f3443wtckeqavwp9950wacxfmwv20q'):
        
        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.clients = {}
        self.job_counter = 0
        self.current_job = None
        self.running = True
        self.extranonce1_counter = 0
        
        print(f"🔥 RinCoin PRODUCTION Stratum Proxy Server")
        print(f"Stratum: {stratum_host}:{stratum_port}")
        print(f"RPC: {rpc_host}:{rpc_port}")
        print(f"Target: {target_address}")
        
    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=10)
            
            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 create_coinbase_tx(self, template, extranonce1, extranonce2):
        """Create coinbase transaction"""
        try:
            # Get coinbase value (block reward + fees)
            coinbase_value = template.get('coinbasevalue', 2500000000)  # 25 RIN in satoshis
            
            # Create coinbase transaction
            # Version (4 bytes)
            coinbase_tx = struct.pack('<L', 1)
            
            # Input count (1 byte) - always 1 for coinbase
            coinbase_tx += b'\x01'
            
            # Previous output hash (32 bytes of zeros for coinbase)
            coinbase_tx += b'\x00' * 32
            
            # Previous output index (4 bytes, 0xffffffff for coinbase)
            coinbase_tx += b'\xff\xff\xff\xff'
            
            # Script length and coinbase script
            height = template.get('height', 0)
            height_bytes = struct.pack('<L', height)[:3]  # BIP34 - height in coinbase
            
            # Coinbase script: height + extranonces + arbitrary data
            coinbase_script = height_bytes + extranonce1.encode() + extranonce2.encode()
            coinbase_script += b'/RinCoin Stratum Pool/'  # Pool signature
            
            # Script length (varint) + script
            coinbase_tx += struct.pack('B', len(coinbase_script)) + coinbase_script
            
            # Sequence (4 bytes)
            coinbase_tx += b'\xff\xff\xff\xff'
            
            # Output count (1 byte) - 1 output to our address
            coinbase_tx += b'\x01'
            
            # Output value (8 bytes)
            coinbase_tx += struct.pack('<Q', coinbase_value)
            
            # Output script (simplified - you'd need proper address decoding)
            # For now, we'll use a simplified P2WPKH script
            script_pubkey = self.address_to_script_pubkey(self.target_address)
            coinbase_tx += struct.pack('B', len(script_pubkey)) + script_pubkey
            
            # Lock time (4 bytes)
            coinbase_tx += struct.pack('<L', 0)
            
            return coinbase_tx
            
        except Exception as e:
            print(f"Coinbase creation error: {e}")
            return None

    def address_to_script_pubkey(self, address):
        """Convert bech32 address to script pubkey (simplified)"""
        # This is a simplified version - in production you'd use proper bech32 decoding
        # For now, return a standard P2WPKH template
        # TODO: Implement proper bech32 decoding
        return b'\x00\x14' + b'\x00' * 20  # OP_0 + 20-byte pubkey hash placeholder

    def calculate_merkle_root(self, transactions):
        """Calculate merkle root from list of transaction hashes"""
        if not transactions:
            return b'\x00' * 32
            
        # Convert hex strings to bytes
        tx_hashes = [bytes.fromhex(tx) if isinstance(tx, str) else tx for tx in transactions]
        
        while len(tx_hashes) > 1:
            if len(tx_hashes) % 2 == 1:
                tx_hashes.append(tx_hashes[-1])  # Duplicate last hash if odd number
                
            new_level = []
            for i in range(0, len(tx_hashes), 2):
                combined = tx_hashes[i] + tx_hashes[i + 1]
                hash_result = hashlib.sha256(hashlib.sha256(combined).digest()).digest()
                new_level.append(hash_result)
            tx_hashes = new_level
            
        return tx_hashes[0] if tx_hashes else b'\x00' * 32

    def get_block_template(self):
        """Get new block template from RinCoin node"""
        try:
            template = self.rpc_call("getblocktemplate", [{"rules": ["segwit"]}])
            if not template:
                return None
                
            self.job_counter += 1
            
            # Calculate target from bits
            bits = template.get('bits', '1d00ffff')
            target = self.bits_to_target(bits)
            
            # Prepare transaction list (without coinbase)
            transactions = template.get('transactions', [])
            tx_hashes = [bytes.fromhex(tx['hash'])[::-1] for tx in transactions]  # Reverse for little-endian
            
            job = {
                "job_id": f"job_{self.job_counter:08x}",
                "template": template,
                "prevhash": template.get("previousblockhash", "0" * 64),
                "version": template.get('version', 1),
                "bits": bits,
                "ntime": int(time.time()),
                "target": target,
                "transactions": transactions,
                "tx_hashes": tx_hashes,
                "height": template.get('height', 0),
                "coinbasevalue": template.get('coinbasevalue', 2500000000)
            }
            
            self.current_job = job
            print(f"📦 New job: {job['job_id']} | Height: {job['height']} | Reward: {job['coinbasevalue']/100000000:.2f} RIN")
            return job
            
        except Exception as e:
            print(f"Get block template error: {e}")
            return None

    def bits_to_target(self, bits_hex):
        """Convert bits to target (difficulty)"""
        try:
            bits = int(bits_hex, 16)
            exponent = bits >> 24
            mantissa = bits & 0xffffff
            target = mantissa * (256 ** (exponent - 3))
            return f"{target:064x}"
        except:
            return "0000ffff00000000000000000000000000000000000000000000000000000000"

    def construct_block_header(self, job, extranonce1, extranonce2, ntime, nonce):
        """Construct block header for submission"""
        try:
            # Create coinbase transaction
            coinbase_tx = self.create_coinbase_tx(job['template'], extranonce1, extranonce2)
            if not coinbase_tx:
                return None, None
                
            # Calculate coinbase hash
            coinbase_hash = hashlib.sha256(hashlib.sha256(coinbase_tx).digest()).digest()[::-1]  # Reverse for little-endian
            
            # Create full transaction list (coinbase + other transactions)
            all_tx_hashes = [coinbase_hash] + job['tx_hashes']
            
            # Calculate merkle root
            merkle_root = self.calculate_merkle_root(all_tx_hashes)
            
            # Construct block header (80 bytes)
            header = b''
            header += struct.pack('<L', job['version'])  # Version (4 bytes)
            header += bytes.fromhex(job['prevhash'])[::-1]  # Previous block hash (32 bytes, reversed)
            header += merkle_root[::-1]  # Merkle root (32 bytes, reversed)
            header += struct.pack('<L', int(ntime, 16))  # Timestamp (4 bytes)
            header += bytes.fromhex(job['bits'])[::-1]  # Bits (4 bytes, reversed)
            header += struct.pack('<L', int(nonce, 16))  # Nonce (4 bytes)
            
            # Construct full block
            block = header
            
            # Transaction count (varint)
            tx_count = 1 + len(job['transactions'])
            if tx_count < 253:
                block += struct.pack('B', tx_count)
            else:
                block += b'\xfd' + struct.pack('<H', tx_count)
            
            # Add coinbase transaction
            block += coinbase_tx
            
            # Add other transactions
            for tx in job['transactions']:
                block += bytes.fromhex(tx['data'])
            
            return header, block
            
        except Exception as e:
            print(f"Block construction error: {e}")
            return None, None

    def validate_and_submit_block(self, job, extranonce1, extranonce2, ntime, nonce):
        """Validate proof of work and submit block if valid"""
        try:
            # Construct block
            header, full_block = self.construct_block_header(job, extranonce1, extranonce2, ntime, nonce)
            if not header or not full_block:
                return False, "Block construction failed"
            
            # Calculate block hash (double SHA256 of header)
            block_hash = hashlib.sha256(hashlib.sha256(header).digest()).digest()
            
            # Convert to hex (reversed for display)
            block_hash_hex = block_hash[::-1].hex()
            
            # Check if hash meets target (proof of work validation)
            target_int = int(job['target'], 16)
            hash_int = int(block_hash_hex, 16)
            
            print(f"🔍 Hash: {block_hash_hex}")
            print(f"🎯 Target: {job['target']}")
            print(f"✅ Valid PoW: {hash_int < target_int}")
            
            if hash_int < target_int:
                # Valid block! Submit to node
                block_hex = full_block.hex()
                
                print(f"🚀 Submitting block {block_hash_hex[:16]}...")
                result = self.rpc_call("submitblock", [block_hex])
                
                if result is None:  # Success
                    print(f"🎉 BLOCK ACCEPTED! Hash: {block_hash_hex}")
                    print(f"💰 Reward: {job['coinbasevalue']/100000000:.2f} RIN -> {self.target_address}")
                    return True, "Block accepted"
                else:
                    print(f"❌ Block rejected: {result}")
                    return False, f"Block rejected: {result}"
            else:
                # Valid share but not a block
                return True, "Share accepted"
                
        except Exception as e:
            print(f"Block 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"ex{self.extranonce1_counter:06x}"
                
                # 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
                ])
                
                # Send difficulty (simplified - always 1 for now)
                self.send_stratum_notification(client, "mining.set_difficulty", [1])
                
                # Send initial job
                if self.current_job:
                    self.send_job_to_client(client, self.current_job)
                else:
                    # Get new job if none exists
                    if self.get_block_template():
                        self.send_job_to_client(client, self.current_job)
            
            elif method == "mining.authorize":
                # Authorization (accept any user/pass for now)
                username = params[0] if params else "anonymous"
                self.clients[addr]['username'] = username
                self.send_stratum_response(client, msg_id, True)
                print(f"[{addr}] Authorized as {username}")
            
            elif method == "mining.submit":
                # Submit share/block
                if len(params) >= 5:
                    username = params[0]
                    job_id = params[1]
                    extranonce2 = params[2]
                    ntime = params[3]
                    nonce = params[4]
                    
                    print(f"[{addr}] Submit: job={job_id}, nonce={nonce}")
                    
                    # Validate submission
                    if self.current_job and job_id == self.current_job['job_id']:
                        extranonce1 = self.clients[addr].get('extranonce1', 'ex000000')
                        
                        # Validate and potentially submit block
                        success, message = self.validate_and_submit_block(
                            self.current_job, extranonce1, extranonce2, ntime, nonce
                        )
                        
                        if success:
                            self.send_stratum_response(client, msg_id, True)
                            if "Block accepted" in message:
                                # Broadcast new job after block found
                                threading.Thread(target=self.update_job_after_block, daemon=True).start()
                        else:
                            self.send_stratum_response(client, msg_id, False, message)
                    else:
                        self.send_stratum_response(client, msg_id, False, "Stale job")
                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"""
        try:
            self.send_stratum_notification(client, "mining.notify", [
                job["job_id"],
                job["prevhash"],
                "",  # coinb1 (empty - we handle coinbase internally)
                "",  # coinb2 (empty - we handle coinbase internally)  
                [],  # merkle_branch (empty - we calculate merkle root)
                f"{job['version']:08x}",
                job["bits"],
                f"{job['ntime']:08x}",
                True  # clean_jobs
            ])
        except Exception as e:
            print(f"Failed to send job: {e}")

    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 job {self.current_job['job_id']} 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"""
        print(f"[{addr}] Connected")
        if addr not in self.clients:
            self.clients[addr] = {}
        self.clients[addr]['socket'] = client
        
        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]
            print(f"[{addr}] Disconnected")

    def job_updater(self):
        """Periodically update mining jobs"""
        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()
                        
            except Exception as e:
                print(f"Job updater error: {e}")

    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")
            print(f"📊 Current height: {blockchain_info.get('blocks', 'unknown')}")
            print(f"⛓️  Chain: {blockchain_info.get('chain', 'unknown')}")
            
            # 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)
            
            print(f"🚀 PRODUCTION 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']}")
            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("\n🛑 Shutting 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__":
    proxy = RinCoinStratumProxy()
    proxy.start()