import csv

from pyspark.sql import types as T
from transforms.api import Input, Output, configure, transform
from transforms.sidecar import Volume, sidecar
from myproject.datasets.sidecar_client import copy_artifacts, execute_cli

FIELDS = ("data1", "data2", "data3")


@configure(profile=["EXECUTOR_CORES_EXTRA_SMALL", "NUM_EXECUTORS_4"])
@sidecar(image="simple-cli", tag="0.0.1", volumes=[Volume("shared")])
@transform(
    output=Output("/training/containers/cli-parallel-files"),
    output_rows=Output("/training/containers/cli-parallel-rows"),
    source=Input("/training/containers/input"),
)
def compute(output, output_rows, source, ctx):
    source_frame = source.dataframe()

    def process_row(row):
        def prepare_input(data_directory):
            with (data_directory / "infile.csv").open("w", newline="", encoding="utf-8") as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(FIELDS)
                writer.writerow([row[name] for name in FIELDS])

        with execute_cli(prepare_input) as request:
            with (request / "data" / "outfile.csv").open(newline="", encoding="utf-8") as csvfile:
                reader = csv.DictReader(csvfile)
                if reader.fieldnames != list(FIELDS):
                    raise ValueError("Expected the CSV columns data1, data2, data3")
                values = next(reader, None)
                if values is None or next(reader, None) is not None:
                    raise ValueError("Expected exactly one output CSV row")
                if set(values) != set(FIELDS) or any(values[name] is None for name in FIELDS):
                    raise ValueError("Expected three values in the output CSV row")
            copy_artifacts(output, request, prefix=request.name + "_", include_input=True)
            return tuple(row[name] for name in FIELDS) + ("success",) + tuple(values[name] for name in FIELDS)

    schema = T.StructType(
        [source_frame.schema[name] for name in FIELDS]
        + [T.StructField("success", T.StringType(), False)]
        + [T.StructField("out" + name, T.StringType(), False) for name in FIELDS]
    )
    rows = source_frame.select(*FIELDS).repartition(4).rdd.map(process_row)
    output_rows.write_dataframe(ctx.spark_session.createDataFrame(rows, schema))