import socket
import time

import requests
from pyspark.sql import types as T
from transforms.api import Output, transform_df
from transforms.sidecar import sidecar


@sidecar(image="simple-http", tag="0.0.1")
@transform_df(Output("/training/containers/http-response"))
def compute(ctx):
    def call_sidecar(_):
        deadline = time.monotonic() + 60
        while True:
            try:
                with socket.create_connection(("127.0.0.1", 1234), timeout=1):
                    break
            except OSError as error:
                if time.monotonic() >= deadline:
                    raise TimeoutError("Sidecar HTTP port did not become ready") from error
                time.sleep(0.1)
        response = requests.get("http://127.0.0.1:1234/hello", timeout=(5, 30))
        response.raise_for_status()
        return (response.text,)

    rows = ctx.spark_session.sparkContext.parallelize([1], 1).map(call_sidecar)
    schema = T.StructType([T.StructField("response_text", T.StringType(), False)])
    return ctx.spark_session.createDataFrame(rows, schema)