from dagster import op, job, In, Out
import pysera
from api.execution_context import ExecutionContext
from api.assets import Asset


@op(out={"image_input": Out(str), "mask_input": Out(str)})
def image_reader_op():
    image_input = "data/sample/images"
    mask_input  = "data/sample/masks"
    return {"image_input": image_input, "mask_input": mask_input}


@op(ins={"image_input": In(str), "mask_input": In(str)}, out={"features": Out(Asset)})
def pysera_extractor_op(image_input, mask_input):
    output_dir = "results"

    result = pysera.process_batch(
        image_input=image_input,
        mask_input=mask_input,
        output_path=output_dir,
        num_workers="auto",
        enable_parallelism=True,
        apply_preprocessing=True,
        categories="all",
        dimensions="1st,2_5d,3d",
        feature_value_mode="REAL_VALUE",
        extraction_mode="handcrafted_feature",
        report="info",
    )

    df = result.get("features_extracted")
    return {"features": Asset(df, "pysera_features_df")}

@op(ins={"features": In(Asset)}, out={"csv_path": Out(Asset)})
def csv_writer_op(features):
    import os
    import csv
    from datetime import datetime

    df = features.data
    rows = []
    if hasattr(df, "iterrows"):
        for _, row in df.iterrows():
            rows.append({"name": row[0], "value": row[1]})

    os.makedirs("results", exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    path = os.path.join("results", f"features_dagster_{ts}.csv")

    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["name", "value"])
        for r in rows:
            w.writerow([r["name"], r["value"]])

    return {"csv_path": Asset(path, "csv_path")}

@job
def radiuma_pysera_job():
    image_mask = image_reader_op()
    feats = pysera_extractor_op(
        image_input=image_mask["image_input"],
        mask_input=image_mask["mask_input"],
    )
    csv_writer_op(features=feats["features"])
