import socket

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

BUFSIZE = int(1024**2)

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

def receive(ip, port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect((ip, port))
        log("Connected! Starting to receive data")
        total_bytes = 0
        while True:
            data = sock.recv(BUFSIZE)
            if not data:
                break
            total_bytes += len(data)
        log(f"Received {total_bytes // 1024**2}MB in total")

@trigger(event='decoder_ready')
class StreamTrainer(FlowSpec):

    ip = Parameter('ip', default='no_ip')
    port = Parameter('port', default=0)

    @card(type='blank', refresh_interval=1)
    @step
    def start(self):
        log(f"🔌 Connecting to {self.ip}:{self.port}")
        receive(self.ip, self.port)
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
    StreamTrainer()