from dagster import (
    AssetKey,
    AssetMaterialization,
    DagsterEventType,
    DagsterRunStatus,
    EventRecordsFilter,
    Field,
    OpExecutionContext,
    RunsFilter,
    op,
)

@op(
    config_schema={
        "asset_key": Field(
            config=str,
            description="The name of the asset that the job materializes",
        ),
        "job_name": Field(
            config=str,
            description="Name of the job (in this repo) to perform a backfill for",
        ),
        "repository_name": Field(
            config=str,
            description="Name of the dagster repository",
        ),
    },
)
def partitioned_asset_backfill_op(context: OpExecutionContext):
    asset_key = AssetKey.from_user_string(context.op_config["asset_key"])

materializations = context.instance.get_event_records(
        EventRecordsFilter(
            event_type=DagsterEventType.ASSET_MATERIALIZATION,
            asset_key=asset_key,
        ),
    )

job_partition_history = context.instance.get_run_partition_data(
        runs_filter=RunsFilter(
            job_name=context.op_config["job_name"],
            tags={".dagster/repository": context.op_config["repository_name"]},
            statuses=[DagsterRunStatus.SUCCESS],
        )
    )

asset_materialized_partitions = [
        materialization.event_log_entry.dagster_event.partition for materialization in materializations
    ]

job_success_partitions = [run.partition for run in job_partition_history]
    missing_partitions = set(job_success_partitions) - set(asset_materialized_partitions)
    for partition in missing_partitions:
        context.log_event(
            AssetMaterialization(
                asset_key=asset_key,
                partition=partition,
                description="Asset previously materialized as a job/op",
            )
        )