import socket
import time

from metaflow import FlowSpec, step, card, current, Parameter
from metaflow.cards import ProgressBar
from metaflow.cards import Markdown as MD
from metaflow.integrations import ArgoEvent

PORT = 13200

def detect_primary_ip():
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]

def open_port(host='0.0.0.0', port=PORT):
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((host, port))
    srv.listen(1)
    return srv

def send_data(sock, amount):
    one_mb = b'a' * int(1024**2)
    num_mbs = amount * 1000
    p = ProgressBar(max=num_mbs, label='MBs sent')
    current.card.append(p)
    t = time.time()
    for i in range(num_mbs):
        to_send = len(one_mb)
        sent = 0
        while to_send > 0:
            n = sock.send(one_mb[sent:])
            to_send -= n
        p.update(i)
        current.card.refresh()
    elapsed = time.time() - t
    log("⏩ Sent %2.2f GBit/s" % (8 * amount / elapsed))

def log(msg):
    current.card.append(MD(f"## {msg}"))
    current.card.refresh()

class DecodingFlow(FlowSpec):

    amount = Parameter('amount', default=10, help='number of GBs to send')

    @card(type='blank', refresh_interval=1)
    @step
    def start(self):
        current.card.refresh()
        my_ip = detect_primary_ip()
        log(f"🔌 IP {my_ip}")
        srv = open_port()
        log("Ready for serving, triggering `decoding_ready` event")
        ArgoEvent('decoder_ready').publish(payload={'ip': my_ip, 'port': PORT})
        log("Waiting for incoming connection...")
        client_sock, client_addr = srv.accept()
        log(f"Connected from {client_addr[0]}:{client_addr[1]}")
        send_data(client_sock, self.amount)
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
    DecodingFlow()