from dagster import EventMetadataEntry, check, seven
from dagster.core.definitions import (
    AssetMaterialization,
    ExpectationResult,
    Materialization,
    SolidHandle,
)
from dagster.core.definitions.events import PythonArtifactMetadataEntryData
from dagster.core.events import (
    DagsterEvent,
    DagsterEventType,
    EngineEventData,
    StepExpectationResultData,
    StepMaterializationData,
)
from dagster.core.execution.plan.objects import (
    StepFailureData,
    StepInputData,
    StepOutputData,
    StepOutputHandle,
    StepRetryData,
    StepSuccessData,
    TypeCheckData,
    UserFailureData,
)
from dagster.utils.error import SerializableErrorInfo

HANDLED_EVENTS = {
    "ExecutionStepStartEvent": DagsterEventType.STEP_START,
    "ExecutionStepInputEvent": DagsterEventType.STEP_INPUT,
    "ExecutionStepOutputEvent": DagsterEventType.STEP_OUTPUT,
    "ExecutionStepFailureEvent": DagsterEventType.STEP_FAILURE,
    "ExecutionStepSkippedEvent": DagsterEventType.STEP_SKIPPED,
    "ExecutionStepSuccessEvent": DagsterEventType.STEP_SUCCESS,
    "StepMaterializationEvent": DagsterEventType.STEP_MATERIALIZATION,
    "StepExpectationResultEvent": DagsterEventType.STEP_EXPECTATION_RESULT,
    "ObjectStoreOperationEvent": DagsterEventType.OBJECT_STORE_OPERATION,
    "EngineEvent": DagsterEventType.ENGINE_EVENT,
    "ExecutionStepUpForRetryEvent": DagsterEventType.STEP_UP_FOR_RETRY,
    "ExecutionStepRestartEvent": DagsterEventType.STEP_RESTARTED,
}


def expectation_result_from_data(data):
    return ExpectationResult(
        success=data["success"],
        label=data["label"],
        description=data.get("description"),  # enforce?
        metadata_entries=list(event_metadata_entries(data.get("metadataEntries")) or []),
    )


def materialization_from_data(data):
    if data.get("asset_key"):
        return AssetMaterialization(
            asset_key=data["asset_key"],
            description=data.get("description"),
            metadata_entries=list(event_metadata_entries(data.get("metadataEntries")) or []),
        )
    return Materialization(
        label=data["label"],
        description=data.get("description"),  # enforce?
        metadata_entries=list(event_metadata_entries(data.get("metadataEntries")) or []),
    )


def error_from_data(data):
    return SerializableErrorInfo(
        message=data["message"],
        stack=data["stack"],
        cls_name=data["className"],
        cause=error_from_data(data["cause"]) if data.get("cause") else None,
    )


def event_metadata_entries(metadata_entry_datas):
    if not metadata_entry_datas:
        return

    for metadata_entry_data in metadata_entry_datas:
        typename = metadata_entry_data["__typename"]
        label = metadata_entry_data["label"]
        description = metadata_entry_data.get("description")
        if typename == "EventPathMetadataEntry":
            yield EventMetadataEntry.path(
                label=label, description=description, path=metadata_entry_data["path"]
            )
        elif typename == "EventJsonMetadataEntry":
            yield EventMetadataEntry.json(
                label=label,
                description=description,
                data=seven.json.loads(metadata_entry_data.get("jsonString", "")),
            )
        elif typename == "EventMarkdownMetadataEntry":
            yield EventMetadataEntry.md(
                label=label, description=description, md_str=metadata_entry_data.get("md_str", "")
            )
        elif typename == "EventTextMetadataEntry":
            yield EventMetadataEntry.text(
                label=label, description=description, text=metadata_entry_data["text"]
            )
        elif typename == "EventUrlMetadataEntry":
            yield EventMetadataEntry.url(
                label=label, description=description, url=metadata_entry_data["url"]
            )
        elif typename == "EventPythonArtifactMetadataEntry":
            yield EventMetadataEntry(
                label=label,
                description=description,
                entry_data=PythonArtifactMetadataEntryData(
                    metadata_entry_data["module"], metadata_entry_data["name"]
                ),
            )
        else:
            check.not_implemented("TODO for type {}".format(typename))


def dagster_event_from_dict(event_dict, pipeline_name):
    check.dict_param(event_dict, "event_dict", key_type=str)
    check.str_param(pipeline_name, "pipeline_name")

    # Get event_type
    event_type = HANDLED_EVENTS.get(event_dict["__typename"])
    if not event_type:
        raise Exception("unhandled event type %s" % event_dict["__typename"])

    # Get event_specific_data
    event_specific_data = None
    if event_type == DagsterEventType.STEP_OUTPUT:
        event_specific_data = StepOutputData(
            step_output_handle=StepOutputHandle(event_dict["stepKey"], event_dict["outputName"]),
            type_check_data=TypeCheckData(
                success=event_dict["typeCheck"]["success"],
                label=event_dict["typeCheck"]["label"],
                description=event_dict.get("description"),
                metadata_entries=list(
                    event_metadata_entries(event_dict.get("metadataEntries")) or []
                ),
            ),
        )

    elif event_type == DagsterEventType.STEP_INPUT:
        event_specific_data = StepInputData(
            input_name=event_dict["inputName"],
            type_check_data=TypeCheckData(
                success=event_dict["typeCheck"]["success"],
                label=event_dict["typeCheck"]["label"],
                description=event_dict.get("description"),
                metadata_entries=list(
                    event_metadata_entries(event_dict.get("metadataEntries")) or []
                ),
            ),
        )
    elif event_type == DagsterEventType.STEP_SUCCESS:
        event_specific_data = StepSuccessData(0.0)

    elif event_type == DagsterEventType.STEP_UP_FOR_RETRY:
        event_specific_data = StepRetryData(
            error=error_from_data(event_dict["retryError"]),
            seconds_to_wait=event_dict["secondsToWait"],
        )

    elif event_type == DagsterEventType.STEP_MATERIALIZATION:
        materialization = event_dict["materialization"]
        event_specific_data = StepMaterializationData(
            materialization=materialization_from_data(materialization)
        )
    elif event_type == DagsterEventType.STEP_EXPECTATION_RESULT:
        expectation_result = expectation_result_from_data(event_dict["expectationResult"])
        event_specific_data = StepExpectationResultData(expectation_result)

    elif event_type == DagsterEventType.STEP_FAILURE:
        event_specific_data = StepFailureData(
            error_from_data(event_dict["error"]),
            UserFailureData(
                label=event_dict["failureMetadata"]["label"],
                description=event_dict["failureMetadata"]["description"],
                metadata_entries=list(
                    event_metadata_entries(event_dict.get("metadataEntries")) or []
                ),
            )
            if event_dict.get("failureMetadata")
            else None,
        )

    elif event_type == DagsterEventType.ENGINE_EVENT:
        event_specific_data = EngineEventData(
            metadata_entries=list(event_metadata_entries(event_dict.get("metadataEntries"))),
            marker_start=event_dict.get("markerStart"),
            marker_end=event_dict.get("markerEnd"),
            error=error_from_data(event_dict["engineError"])
            if event_dict.get("engineError")
            else None,
        )

    return DagsterEvent(
        event_type_value=event_type.value,
        pipeline_name=pipeline_name,
        step_key=event_dict.get("stepKey"),
        solid_handle=SolidHandle.from_string(event_dict["solidHandleID"])
        if event_dict.get("solidHandleID")
        else None,
        # at the time of writing this:
        # * 'COMPUTE` is the only step kind
        # * this code should get deleted in the near future as we move away from
        #   dagster-graphql CLI as what we invoke in dask/k8s/etc.
        step_kind_value="COMPUTE" if event_dict.get("stepKey") else None,
        logging_tags=None,
        event_specific_data=event_specific_data,
    )
