import os
import json
from web3 import Web3
from dotenv import load_dotenv

load_dotenv()

RPC_URL = os.environ.get("PURECHAIN_RPC_URL", "http://3.34.161.207:8545")
PRIVATE_KEY = os.environ.get("PURECHAIN_PRIVATE_KEY")
ADDR_OLD = os.environ.get("PURECHAIN_CONTRACT_ADDRESS_PUREBI_JUN11")
ADDR_NEW = os.environ.get("PURECHAIN_CONTRACT_ADDRESS")

def main():
    if not PRIVATE_KEY or not ADDR_OLD or not ADDR_NEW:
        print("❌ Missing environment variables. Check .env")
        return

    w3 = Web3(Web3.HTTPProvider(RPC_URL))
    if not w3.is_connected():
        print("❌ Web3 not connected.")
        return

    admin_account = w3.eth.account.from_key(PRIVATE_KEY)
    print(f"Admin Account: {admin_account.address}")
    
    # ABI for just the functions we need
    abi_minimal = json.loads('''[
        {"anonymous":false,"inputs":[{"indexed":false,"internalType":"uint256","name":"id","type":"uint256"},{"indexed":false,"internalType":"address","name":"contributor","type":"address"},{"indexed":false,"internalType":"string","name":"phrase","type":"string"},{"indexed":false,"internalType":"string","name":"dialect","type":"string"},{"indexed":false,"internalType":"uint256","name":"timestamp","type":"uint256"}],"name":"EntryProposed","type":"event"},
        {"anonymous":false,"inputs":[{"indexed":false,"internalType":"uint256","name":"id","type":"uint256"}],"name":"EntryVerified","type":"event"},
        {"inputs":[{"internalType":"address","name":"","type":"address"},{"internalType":"string","name":"","type":"string"}],"name":"userXP","outputs":[{"internalType":"uint256","name":"","type":"uint256"}],"stateMutability":"view","type":"function"},
        {"inputs":[{"internalType":"address[]","name":"_users","type":"address[]"},{"internalType":"string[]","name":"_dialects","type":"string[]"},{"internalType":"uint256[]","name":"_xps","type":"uint256[]"}],"name":"batchMigrateXP","outputs":[],"stateMutability":"nonpayable","type":"function"}
    ]''')
    
    contract_old = w3.eth.contract(address=w3.to_checksum_address(ADDR_OLD), abi=abi_minimal)
    contract_new = w3.eth.contract(address=w3.to_checksum_address(ADDR_NEW), abi=abi_minimal)

    # We need to discover users and dialects by parsing EntryProposed events
    print("🔍 Discovering users from old contract events...")
    try:
        events = contract_old.events.EntryProposed.get_logs(fromBlock=0, toBlock='latest')
    except Exception as e:
        print(f"Failed to fetch logs: {e}")
        return

    user_dialects = set()
    for ev in events:
        user = ev.args.contributor
        dialect = ev.args.dialect
        user_dialects.add((user, dialect))
        
    if not user_dialects:
        print("No users found.")
        return
        
    print(f"Found {len(user_dialects)} User/Dialect pairs. Checking XP balances...")
    
    users = []
    dialects = []
    xps = []
    
    for (user, dialect) in user_dialects:
        old_xp = contract_old.functions.userXP(user, dialect).call()
        if old_xp > 0:
            # Check if it was already migrated
            new_xp = contract_new.functions.userXP(user, dialect).call()
            if new_xp < old_xp:
                diff = old_xp - new_xp
                users.append(user)
                dialects.append(dialect)
                xps.append(diff)
                print(f"💡 Found {diff} unmigrated XP for {user} in {dialect}")
                
    if not users:
        print("✅ No new XP needs migration.")
        return
        
    print(f"🚀 Migrating {len(users)} records...")
    
    nonce = w3.eth.get_transaction_count(admin_account.address)
    tx = contract_new.functions.batchMigrateXP(users, dialects, xps).build_transaction({
        'chainId': w3.eth.chain_id,
        'gas': 3000000,
        'gasPrice': w3.eth.gas_price,
        'nonce': nonce,
    })
    
    signed_tx = w3.eth.account.sign_transaction(tx, PRIVATE_KEY)
    tx_hash = w3.eth.send_raw_transaction(signed_tx.raw_transaction)
    
    print(f"⏳ Migration TX sent! Hash: {tx_hash.hex()}")
    receipt = w3.eth.wait_for_transaction_receipt(tx_hash)
    print(f"🎉 Migration complete in block {receipt.blockNumber}!")

if __name__ == "__main__":
    main()