# ruff: noqa: E402

from __future__ import annotations

import argparse
import json
import os
import sys
import traceback
from datetime import datetime, timezone
from typing import Any

import fsspec.implementations.arrow as pfs


hopsfs = pfs.HadoopFileSystem("default", user=os.environ["HADOOP_USER_NAME"])
import hopsworks
from hsfs import engine
from hsfs.constructor import query
from hsfs.core import (
    feature_monitoring_config_engine,
    feature_view_engine,
    kafka_engine,
)
from hsfs.statistics_config import StatisticsConfig
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, expr, max, row_number
from pyspark.sql.types import StructField, StructType, _parse_datatype_string
from pyspark.sql.window import Window


def read_job_conf(path: str) -> dict[Any, Any]:
    """
    The configuration file is passed as path on HopsFS
    The path is a JSON containing different values depending on the op type
    """
    file_name = os.path.basename(path)
    hopsfs.download(path, ".")
    with open(file_name, "r") as f:
        return json.loads(f.read())


def setup_spark() -> SparkSession:
    return SparkSession.builder.enableHiveSupport().getOrCreate()


def get_feature_store_handle(feature_store: str = "") -> hopsworks.hsfs.feature_store:
    project = hopsworks.login()
    return project.get_feature_store(feature_store)


def sort_schema(fg_schema: StructType, csv_df_schema: StructType) -> StructType:
    # The schema order of the fg_schema needs to match the
    # order of the csv_df_schema
    csv_df_schema_indices = [
        csv_df_schema.names.index(field) for field in fg_schema.names
    ]
    fg_schema_sorted = sorted(
        zip(fg_schema.fields, csv_df_schema_indices), key=lambda x: x[1]
    )
    return StructType([f[0] for f in fg_schema_sorted])


def get_fg_spark_df(job_conf: dict[Any, Any], fg_schema: StructType) -> Any:
    data_path = job_conf.pop("data_path")
    data_format = job_conf.pop("data_format")
    data_options = job_conf.pop("data_options")

    csv_df = spark.read.format(data_format).options(**data_options).load(data_path)

    schema = sort_schema(fg_schema, csv_df.schema)

    return (
        spark.read.format(data_format)
        .options(**data_options)
        .schema(schema)
        .load(data_path)
    )


def insert_fg(spark: SparkSession, job_conf: dict[Any, Any]) -> None:
    """
    Insert data into a feature group.
    The data path, feature group name and versions are in the configuration file
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    fg = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])

    schema = StructType(
        [StructField(f.name, _parse_datatype_string(f.type), True) for f in fg.features]
    )

    df = get_fg_spark_df(job_conf, schema)
    fg.insert(df, write_options=job_conf.pop("write_options", {}) or {})


def create_td(job_conf: dict[Any, Any]) -> None:
    # Extract the feature store handle
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    # Extract the query object
    q = query.Query._hopsworks_json(job_conf.pop("query"))

    td = fs.get_training_dataset(name=job_conf["name"], version=job_conf["version"])
    td.insert(
        q,
        overwrite=job_conf.pop("overwrite", False) or False,
        write_options=job_conf.pop("write_options", {}) or {},
    )


def create_fv_td(job_conf: dict[Any, Any]) -> None:
    # Extract the feature store handle
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    fv = fs.get_feature_view(name=job_conf["name"], version=job_conf["version"])
    fv_engine = feature_view_engine.FeatureViewEngine(fv.featurestore_id)

    user_write_options = job_conf.pop("write_options", {}) or {}

    training_helper_columns = user_write_options.get("training_helper_columns")
    primary_keys = user_write_options.get("primary_keys")
    event_time = user_write_options.get("event_time")
    spine = job_conf.pop("spine", None)
    fv_engine._compute_training_dataset(
        feature_view_obj=fv,
        user_write_options=user_write_options,
        primary_keys=primary_keys,
        event_time=event_time,
        training_helper_columns=training_helper_columns,
        training_dataset_version=job_conf["td_version"],
        spine_df=None if spine is None else read_staged_spine(fs, spine),
    )


def read_staged_spine(fs: Any, spine: dict[str, Any]) -> Any:
    """The spine the client staged for this job, read back as a Spark DataFrame in spine order.

    The client wrote the file under its own Resources/.hopsworks_spine/ and the backend checked it
    against the feature view before naming it in the job configuration. The job rebuilds the query
    from the feature view, so without this the training data would silently be the root feature
    group's rows rather than the population the caller asked for. A file that is gone, because the
    sweeper's featurestore_asof_spine_max_file_age_ms passed before the job ran or because it was
    removed by hand, is an error: the job cannot build the population it was given and must not
    build another.
    """
    from hopsworks_common.client.exceptions import FeatureStoreException
    from hsfs.constructor.inference_spine import ROW_ID_COLUMN, SPINE_DIR

    basename = spine.get("parquetBasename")
    if not basename:
        raise FeatureStoreException(
            "The training dataset job was given a spine that names no staged file."
        )
    path = f"hdfs:///Projects/{fs.project_name}/{SPINE_DIR}/{basename}"
    try:
        frame = setup_spark().read.parquet(path)
    except Exception as e:
        raise FeatureStoreException(
            f"The spine this training dataset was built from could not be read at {path}: {e}."
            " Spine files are kept for featurestore_asof_spine_max_file_age_ms after they are"
            " written; recreate the training dataset with the same `spine_df` to stage it again."
        ) from e
    return frame.orderBy(ROW_ID_COLUMN).drop(ROW_ID_COLUMN)


def compute_stats(job_conf: dict[Any, Any]) -> None:
    """
    Compute/Update statistics on a feature group.

    When `end_commit_time` is present in job_conf (set by the backend from the
    POST /compute query params), statistics are scoped to that commit so they
    are persisted against the specific batch rather than over head.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    entity_type = job_conf["type"]
    if entity_type == "fg":
        entity = fs.get_feature_group(
            name=job_conf["name"], version=job_conf["version"]
        )
    elif entity_type == "external_fg":
        entity = fs.get_external_feature_group(
            name=job_conf["name"], version=job_conf["version"]
        )
    else:
        fv = fs.get_feature_view(job_conf["name"], version=job_conf["version"])
        entity = fv._feature_view_engine._get_training_dataset_metadata(
            feature_view_obj=fv,
            training_dataset_version=job_conf["td_version"],
        )

    end_commit_time = job_conf.get("end_commit_time")
    if end_commit_time is not None and entity_type == "fg":
        # Commit-scoped: read the FG as-of the commit and persist stats against it.
        entity._statistics_engine.compute_and_save_statistics(
            entity, feature_group_commit_id=int(end_commit_time)
        )
    else:
        entity.compute_statistics()


def ge_validate(job_conf: dict[Any, Any]) -> None:
    """
    Run expectation suite attached to a feature group.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    entity = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])

    # when user runs job we always want to save the report and actually perform validation,
    # no matter of setting on feature group level
    entity.validate(
        dataframe=None, save_report=True, validation_options={"run_validation": True}
    )


def import_fg(job_conf: dict[Any, Any]) -> None:
    """
    Import data to a feature group using storage connector.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)
    # retrieve connector
    st = fs.get_storage_connector(name=job_conf["storageConnectorName"])
    # first read data from connector
    spark_options = job_conf.pop("options")
    df = st.read(query=(job_conf.pop("query", "") or ""), options=spark_options)
    # store dataframe into feature group
    if job_conf["statisticsConfig"]:
        stat_config = StatisticsConfig.from_response_json(job_conf["statisticsConfig"])
    else:
        stat_config = None
    # create fg and insert
    fg = fs.get_or_create_feature_group(
        name=job_conf["featureGroupName"],
        version=job_conf["version"],
        primary_key=job_conf["primaryKey"],
        online_enabled=job_conf.pop("onlineEnabled", False) or False,
        statistics_config=stat_config,
        partition_key=job_conf.pop("partitionKey", []) or [],
        description=job_conf["description"],
        event_time=job_conf.pop("eventTime", None) or None,
    )
    fg.insert(df)


def run_feature_monitoring(
    job_conf: dict[str, str], end_commit_time: int | None = None
) -> None:
    """
    Run feature monitoring for a given entity (feature_group or feature_view)
    based on a feature monitoring configuration.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    if job_conf["entity_type"].upper() == "FEATUREGROUPS":
        entity = fs.get_feature_group(
            name=job_conf["name"], version=job_conf["version"]
        )
        feature_group_id = entity._id
        feature_view_name, feature_view_version = None, None
    else:
        feature_group_id = None
        entity = fs.get_feature_view(name=job_conf["name"], version=job_conf["version"])
        feature_view_name, feature_view_version = (
            entity.name,
            entity.version,
        )

    monitoring_config_engine = (
        feature_monitoring_config_engine.FeatureMonitoringConfigEngine(
            feature_store_id=fs._id,
            feature_group_id=feature_group_id,
            feature_view_name=feature_view_name,
            feature_view_version=feature_view_version,
        )
    )

    try:
        monitoring_config_engine._run_feature_monitoring(
            entity=entity,
            config_name=job_conf["config_name"],
            end_commit_time=end_commit_time,
        )
    except Exception as e:
        config = monitoring_config_engine._get_feature_monitoring_configs(
            name=job_conf["config_name"]
        )
        monitoring_config_engine._result_engine._save_with_exception(
            feature_monitoring_config_id=config.id,
            job_name=config.job_name,
        )
        raise e


def delta_vacuum_fg(spark: SparkSession, job_conf: dict[Any, Any]) -> None:
    """
    Run delta vacuum on a feature group.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    entity = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])

    entity.delta_vacuum()


def offline_fg_materialization(
    spark: SparkSession, job_conf: dict[Any, Any], initial_check_point_string: str
) -> None:
    """
    Run materialization job on a feature group.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    entity = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])

    read_options = kafka_engine._get_kafka_config(
        entity.feature_store_id, {}, engine="spark"
    )

    # get starting offsets
    location = entity.prepare_spark_location()
    offset_location = location + "/kafka_offsets"
    # The offsets a run intends to consume, written before its append and removed after the
    # offsets file is saved. A run that finds one repeats exactly that range, so its Delta
    # transaction version matches the earlier attempt's and an append that did commit is skipped
    # instead of being widened by rows that arrived since.
    # It lives beside the feature group directory, not inside it like kafka_offsets: on the
    # first run the table does not exist yet, and Delta refuses to create one at a location
    # that already holds a file (DELTA_MISSING_DELTA_TABLE), which left every later run of a
    # clustered feature group failing on the file the first one wrote. Deleting the feature
    # group removes only its directory, so the name carries the group's id: a group recreated
    # under the same name and version has a new id and never reads the old group's claim.
    pending_offset_location = (
        f"{location.rstrip('/')}_kafka_offsets_pending_{entity.id}"
    )
    pending_offsets = _pending_offsets(
        spark, location, offset_location + "_pending", pending_offset_location
    )
    try:
        if initial_check_point_string:
            starting_offset_string = json.dumps(
                _build_offsets(initial_check_point_string)
            )
        else:
            starting_offset_string = spark.read.json(offset_location).toJSON().first()
    except Exception as e:
        print(f"Failed to use existing offsets: {e}")
        starting_offset_string = None

    # get the current low watermark offsets for all partitions
    low_offsets_string = kafka_engine._kafka_get_offsets(
        topic_name=entity._online_topic_name,
        feature_store_id=entity.feature_store_id,
        offline_write_options={},
        high=False,
    )
    low_offsets = _build_offsets(low_offsets_string)

    if not starting_offset_string:
        # Nothing said where this run should start, so the low watermark would be the whole
        # retained topic. Bound it by when the feature group came into existence instead.
        low_offsets = _offsets_since_creation(
            entity, low_offsets, write_options_of(job_conf)
        )

    # validate and reconcile saved offsets against current topic state
    starting_offset_string = json.dumps(
        _reconcile_offsets(
            starting_offset_string, low_offsets, entity._online_topic_name
        )
    )
    print(f"startingOffsets: {starting_offset_string}")

    # get ending offsets
    ending_offset_string = kafka_engine._kafka_get_offsets(
        topic_name=entity._online_topic_name,
        feature_store_id=entity.feature_store_id,
        offline_write_options={},
        high=True,
    )
    ending_offset_string = json.dumps(_build_offsets(ending_offset_string))
    appends = write_options_of(job_conf).get("operation") == "insert"
    if pending_offsets and appends:
        ending_offset_string = pending_offsets
        print(
            f"repeating the range of an unfinished run, endingOffsets: {ending_offset_string}"
        )
    print(f"endingOffsets: {ending_offset_string}")
    if appends:
        # Claim the range before reading it, not after the append. Written after,
        # it fences nothing: a second execution starting in between (a manual
        # materialize_log(), or the self re-trigger) reads a wider range, derives
        # a different txnVersion and appends the overlap a second time.
        spark.createDataFrame([json.loads(ending_offset_string)]).coalesce(
            1
        ).write.mode("overwrite").json(pending_offset_location)

    # read kafka topic
    df = (
        spark.read.format("kafka")
        .options(**read_options)
        .option("subscribe", entity._online_topic_name)
        .option("startingOffsets", starting_offset_string)
        .option("endingOffsets", ending_offset_string)
        .option("includeHeaders", "true")
        .option("failOnDataLoss", "false")
        .load()
    )

    # filter only the necassary entries
    filtered_df = df.filter(
        expr(
            "CAST(filter(headers, header -> header.key = 'featureGroupId')[0].value AS STRING)"
        )
        == str(entity._id)
    )
    filtered_df = filtered_df.filter(
        expr(
            "CAST(filter(headers, header -> header.key = 'subjectId')[0].value AS STRING)"
        )
        == str(entity.subject["id"])
    )

    # The storage header names the consumer a record is meant for: "online" is for OnlineFS
    # alone, because whoever produced it already wrote the offline leg itself (or has none),
    # so materializing it would write the same rows to the table a second time.
    # Anything else belongs here: the header is absent when both consumers read the record,
    # "offline" when this job is the only one that should, and "1"/"0" from clients that
    # predate this contract meant "ingest online"/"skip online" without ever excluding the
    # offline table.
    storage_header = expr(
        "CAST(filter(headers, header -> header.key = 'storage')[0].value AS STRING)"
    )
    filtered_df = filtered_df.filter(
        storage_header.isNull() | (storage_header != "online")
    )

    # A delete tombstone is on the topic for OnlineFS only: remove_rows has already applied the
    # offline delete directly to the table, so materializing the tombstone re-inserts the key it
    # deleted.
    # New clients mark it "online" above, so this covers the ones that do not.
    # The header is absent on inserts, so null means keep.
    # Exact match on the value, as in OnlineFsHandler.getRow.
    operation_header = expr(
        "CAST(filter(headers, header -> header.key = 'operation')[0].value AS STRING)"
    )
    filtered_df = filtered_df.filter(
        operation_header.isNull() | (operation_header != "delete")
    )

    # limit the number of records ingested
    # default limit is 5M
    limit = 5000000
    write_options = write_options_of(job_conf)
    limit = int(write_options.get("job_limit", limit))
    filtered_df = filtered_df.limit(limit)

    # deserialize dataframe so that it can be properly saved
    deserialized_df = engine._get_instance()._deserialize_from_avro(entity, filtered_df)

    # de-duplicate records
    # timestamp cannot be relied on to order the records in case of duplicates, if they are produced together they would have the same timestamp.
    # Instead use offset to order the records, they are strictly increasing within a partition and since we use primary keys for generating Kafka message keys duplicates are guaranteed to be in the same partition.
    if entity.primary_key:
        partition_columns = [f"value.{key}" for key in entity.primary_key]
        if entity.event_time:
            partition_columns.append(f"value.{entity.event_time}")
        if entity.partition_key:
            partition_columns.extend([f"value.{key}" for key in entity.partition_key])
        window = Window.partitionBy(partition_columns).orderBy(col("offset").desc())
        deduped_df = (
            deserialized_df.withColumn("row_num", row_number().over(window))
            .filter("row_num = 1")
            .drop("row_num")
        )
    else:
        deduped_df = deserialized_df

    # get only the feature values (remove kafka metadata)
    deduped_df = deduped_df.select("value.*")

    # get offsets (do it before inserting to avoid skipping records if data was deleted during the job execution)
    filtered_count = filtered_df.count()
    df_offsets = (
        (df if limit > filtered_count else filtered_df)
        .groupBy("partition")
        .agg(max("offset").alias("offset"))
        .collect()
    )
    offset_dict = json.loads(starting_offset_string)
    for offset_row in df_offsets:
        offset_dict[f"{entity._online_topic_name}"][f"{offset_row.partition}"] = (
            offset_row.offset + 1
        )
    # insert data
    entity.stream = False  # to make sure we dont write to kafka

    # Do not apply transformation function at this point since the data written to Kafka already has transformations applied.
    # A feature log is append-only, so its job configuration carries operation=insert; every other
    # stream group keeps the upsert on its primary key. An append has no key to make a retry
    # idempotent, so the run identifies itself to Delta by the offsets it consumed: a rerun after a
    # commit that never reached the offset file carries the same version and Delta skips it.
    if write_options.get("operation") == "insert":
        write_options = dict(write_options)
        write_options["txnAppId"] = f"hopsworks_feature_log_materialization_{entity.id}"
        write_options["txnVersion"] = str(
            sum(int(v) for v in offset_dict[f"{entity._online_topic_name}"].values())
        )
    if filtered_count == 0 and _offline_table_exists(spark, entity, location):
        print("No records for the offline table in this range, skipping the insert")
    else:
        entity.insert(
            deduped_df,
            storage="offline",
            operation=write_options.get("operation", "upsert"),
            transform=False,
            write_options=write_options,
            validation_options={"schema_validation": False},
        )

    # save offsets
    offset_df = spark.createDataFrame([offset_dict])
    offset_df.coalesce(1).write.mode("overwrite").json(offset_location)
    if write_options.get("operation") == "insert":
        _remove_path(spark, pending_offset_location)


def update_table_schema_fg(spark: SparkSession, job_conf: dict[Any, Any]) -> None:
    """
    Run table schema update job on a feature group.
    """
    feature_store = job_conf.pop("feature_store")
    fs = get_feature_store_handle(feature_store)

    entity = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])

    entity.stream = False
    engine._get_instance()._update_table_schema(entity)


def write_options_of(job_conf) -> dict:
    return job_conf.get("write_options", {}) or {}


def _path_exists(spark, location: str) -> bool:
    jvm = spark._jvm
    path = jvm.org.apache.hadoop.fs.Path(location)
    return path.getFileSystem(spark._jsc.hadoopConfiguration()).exists(path)


def _offline_table_exists(spark, entity, location: str) -> bool:
    """Whether the offline table exists; until it does, the insert creates it."""
    marker = {"DELTA": "_delta_log", "HUDI": ".hoodie"}.get(
        (entity.time_travel_format or "").upper()
    )
    return marker is not None and _path_exists(
        spark, f"{location.rstrip('/')}/{marker}"
    )


def _remove_path(spark, location: str) -> None:
    jvm = spark._jvm
    path = jvm.org.apache.hadoop.fs.Path(location)
    path.getFileSystem(spark._jsc.hadoopConfiguration()).delete(path, True)


def _move_path(spark, source: str, target: str) -> None:
    jvm = spark._jvm
    src = jvm.org.apache.hadoop.fs.Path(source)
    if not src.getFileSystem(spark._jsc.hadoopConfiguration()).rename(
        src, jvm.org.apache.hadoop.fs.Path(target)
    ):
        raise OSError(f"Could not move {source} to {target}")


def _pending_offsets(
    spark, location: str, legacy_location: str, pending_location: str
) -> str | None:
    """The range an unfinished append claimed, as the JSON offsets string, or None.

    `legacy_location` is where clients before this one wrote the claim: inside the
    table directory, where a claim written before the table was created stopped Delta
    from ever creating it. With no `_delta_log` there, nothing was appended, so the claim
    covers nothing and only stands in the way of the create; it is removed. With a table
    there, the claim belongs to an interrupted append and is moved to `pending_location`
    so the retry repeats its range. A claim in the new place always wins over a legacy
    one, because the two share their starting offsets and the newer range covers the
    older one.

    Absence is the only reading of a missing file. A storage, permission or corruption
    error has to stop the run: treating it as "no pending range" would widen the range
    to the offsets that arrived since and append rows the earlier attempt's transaction
    version no longer covers.
    """
    if _path_exists(spark, legacy_location):
        if _path_exists(spark, pending_location) or not _path_exists(
            spark, location + "/_delta_log"
        ):
            _remove_path(spark, legacy_location)
        else:
            _move_path(spark, legacy_location, pending_location)
    if not _path_exists(spark, pending_location):
        return None
    return spark.read.json(pending_location).toJSON().first()


def _timestamp_ms(date_string: str) -> int:
    """Milliseconds since the epoch for a date the backend serialized.

    Parsed here rather than through `hsfs.util`, deliberately. This file ships in the
    spark-feature-pipeline image beside whatever SDK that image happens to carry, so a
    parser reached through the SDK is a second thing that has to be current for the floor
    below to work at all — and when it is not, the floor does not fail loudly, it quietly
    falls back to reading the whole topic.

    `fromisoformat` also accepts the whole of the backend's `yyyy-MM-dd'T'HH:mm:ss.SSSXXX`
    format, including the `+02:00` rendering of that trailing `XXX`, which the SDK's
    pattern list does not.
    """
    parsed = datetime.fromisoformat(date_string.replace("Z", "+00:00"))
    if parsed.tzinfo is None:
        # The backend always sends an offset; anything that does not is read as UTC
        # rather than as the job container's local time.
        parsed = parsed.replace(tzinfo=timezone.utc)
    return int(parsed.timestamp() * 1000)


def _offsets_since_creation(entity, low_offsets: dict, write_options: dict) -> dict:
    """Offsets a first materialization run should start from, floored at the feature group's creation.

    A run gets told where to start by `-initialCheckPointString`, which only the Python
    client passes, or by the offsets file a previous run saved. With neither, the start
    falls back to the topic's low watermark, and the topic is by default
    `<project>_onlinefs`, shared by every online-enabled feature group in the project: the
    run reads the whole retained history of all of them just to drop nearly all of it in
    the featureGroupId filter.

    No record of a feature group can predate the feature group, so its creation time is a
    sound floor. Kafka stamps a record with the producing client's clock by default, so a
    client whose clock lags the backend's could stamp one just before that time;
    `initial_offset_margin_hours` (1 by default) is how far back the floor is moved to
    absorb the skew, and raising it costs only a longer read.

    The low watermark offsets are returned unchanged whenever the floor cannot be
    established, so a failed lookup reads too much rather than too little.
    """
    if not low_offsets or not entity.created:
        return low_offsets

    # Everything the floor is derived from is under the guard: an unparseable margin or
    # creation time is as much a reason to fall back as a broker that will not answer,
    # and none of the three may take the materialization job down with it.
    try:
        margin_hours = float(write_options.get("initial_offset_margin_hours", 1))
        timestamp = _timestamp_ms(entity.created) - int(margin_hours * 60 * 60 * 1000)
        offsets = _build_offsets(
            kafka_engine._kafka_get_offsets_for_times(
                topic_name=entity._online_topic_name,
                feature_store_id=entity.feature_store_id,
                offline_write_options={},
                timestamp=timestamp,
            )
        )
    except Exception as e:
        print(f"Failed to look offsets up by creation time: {e}")
        return low_offsets

    if not offsets:
        return low_offsets

    print(f"No saved offsets, starting from the feature group's creation: {offsets}")
    return offsets


def _build_offsets(initial_check_point_string: str):
    if not initial_check_point_string:
        return ""

    # Split the input string into the topic and partition-offset pairs
    topic, offsets = initial_check_point_string.split(",", 1)
    # Split the offsets and build a dictionary from them
    offsets_dict = {
        partition: int(offset)
        for partition, offset in (pair.split(":") for pair in offsets.split(","))
    }
    # Create the final dictionary structure
    result = {topic: offsets_dict}
    return result


def _reconcile_offsets(
    starting_offset_string: str | None,
    low_offsets: dict,
    topic_name: str,
) -> dict:
    """Reconcile saved offsets against current topic state.

    Handles cases where:
    - No saved offsets exist (returns low watermark offsets)
    - The topic has changed (returns low watermark offsets)
    - Partitions were added (uses low watermark for new partitions)
    - Saved offsets are behind the low watermark (uses low watermark)
    """
    if not starting_offset_string or not low_offsets:
        return low_offsets if low_offsets else {}

    try:
        saved_offsets = json.loads(starting_offset_string)
    except (ValueError, TypeError):
        # Malformed or unexpected JSON — fall back to low watermark offsets
        return low_offsets

    # Expect a mapping of {topic_name: {partition: offset}}
    if not isinstance(saved_offsets, dict):
        return low_offsets

    # topic changed — start from low watermark
    if topic_name not in saved_offsets:
        return low_offsets

    saved_partition_offsets = saved_offsets[topic_name]
    low_partition_offsets = low_offsets.get(topic_name)

    if not isinstance(saved_partition_offsets, dict) or not isinstance(
        low_partition_offsets, dict
    ):
        return low_offsets

    reconciled = {}
    for partition, low_offset in low_partition_offsets.items():
        # Normalize low watermark offset to int; if it cannot be parsed, keep original.
        try:
            low_offset_int = int(low_offset)
        except (TypeError, ValueError):
            reconciled[partition] = low_offset
            continue

        saved_raw_offset = saved_partition_offsets.get(partition)
        try:
            saved_offset_int = (
                int(saved_raw_offset) if saved_raw_offset is not None else None
            )
        except (TypeError, ValueError):
            saved_offset_int = None

        if saved_offset_int is not None and saved_offset_int >= low_offset_int:
            reconciled[partition] = saved_offset_int
        else:
            reconciled[partition] = low_offset_int
    return {topic_name: reconciled}


if __name__ == "__main__":
    # Setup spark first so it fails faster in case of args errors
    # Otherwise the resource manager will wait until the spark application master
    # registers, which never happens.
    spark = setup_spark()

    parser = argparse.ArgumentParser(description="HSFS Job Utils")
    parser.add_argument(
        "-op",
        type=str,
        choices=[
            "insert_fg",
            "create_td",
            "create_fv_td",
            "create_fv_td_spine",
            "compute_stats",
            "ge_validate",
            "import_fg",
            "run_fm",
            "delta_vacuum_fg",
            "offline_fg_materialization",
            "update_table_schema_fg",
        ],
        help="Operation type",
    )
    parser.add_argument(
        "-path",
        type=str,
        help="Location on HopsFS of the JSON containing the full configuration",
    )

    def parse_isoformat_date(da: str) -> datetime:
        # 'Z' is supported in Python 3.11+ so we need to replace it in 3.10
        return datetime.fromisoformat(da.replace("Z", "+00:00"))

    parser.add_argument(
        "-start_time",
        type=parse_isoformat_date,
        help="Job start time",
    )

    parser.add_argument(
        "-initialCheckPointString",
        type=str,
        help="Kafka offset to start consuming from",
    )

    parser.add_argument(
        "-end_commit_time",
        type=int,
        default=None,
        help="Commit timestamp (ms) that triggered this feature monitoring job",
    )
    parser.add_argument(
        "-start_commit_time",
        type=int,
        default=None,
        help=(
            "Optional lower bound commit timestamp (ms). Accepted for wire compatibility "
            "with the statistics/compute endpoint; currently not propagated into the "
            "monitoring window engine for INGESTION configs."
        ),
    )

    args = parser.parse_args()
    job_conf = read_job_conf(args.path)

    success = False
    try:
        if args.op == "insert_fg":
            insert_fg(spark, job_conf)
        elif args.op == "create_td":
            create_td(job_conf)
        elif args.op in ("create_fv_td", "create_fv_td_spine"):
            # The backend launches a spine-anchored job under its own op so an image whose
            # entrypoint predates spines refuses it here, at argument parsing, rather than
            # building the feature view's own rows under the requested version.
            create_fv_td(job_conf)
        elif args.op == "compute_stats":
            compute_stats(job_conf)
        elif args.op == "ge_validate":
            ge_validate(job_conf)
        elif args.op == "import_fg":
            import_fg(job_conf)
        elif args.op == "run_fm":
            run_feature_monitoring(job_conf, end_commit_time=args.end_commit_time)
        elif args.op == "delta_vacuum_fg":
            delta_vacuum_fg(spark, job_conf)
        elif args.op == "offline_fg_materialization":
            offline_fg_materialization(spark, job_conf, args.initialCheckPointString)
        elif args.op == "update_table_schema_fg":
            update_table_schema_fg(spark, job_conf)

        success = True
    except Exception as e:
        print(f"Error : {e}", traceback.format_exc())
    finally:
        if spark is not None:
            try:
                spark.stop()
            except Exception as e:
                print(f"Error stopping spark session: {e}")
        if not success:
            sys.exit(1)
    sys.exit(0)