import math

import graphene
from graphql import (
    DirectiveLocation,
    GraphQLArgument,
    GraphQLDirective,
    GraphQLNonNull,
    GraphQLString,
)

import mlflow
from mlflow.server.graphql.autogenerated_graphql_schema import (
    MlflowExperiment,
    MlflowMetric,
    MlflowModelVersion,
    MlflowRun,
    MlflowSearchRunsInput,
    MlflowSearchRunsResponse,
    MutationType,
    QueryType,
)
from mlflow.utils.proto_json_utils import parse_dict

# Component identifier, to keep compatible with Databricks in-house implementations.
ComponentDirective = GraphQLDirective(
    name="component",
    locations=[
        DirectiveLocation.QUERY,
        DirectiveLocation.MUTATION,
    ],
    args={"name": GraphQLArgument(GraphQLNonNull(GraphQLString))},
)


class Test(graphene.ObjectType):
    output = graphene.String(description="Echoes the input string")


class TestMutation(graphene.ObjectType):
    output = graphene.String(description="Echoes the input string")


class MlflowRunExtension(MlflowRun):
    experiment = graphene.Field(MlflowExperiment)
    model_versions = graphene.List(graphene.NonNull(MlflowModelVersion))

    def resolve_experiment(self, info):
        experiment_id = self.info.experiment_id
        input_dict = {"experiment_id": experiment_id}
        request_message = mlflow.protos.service_pb2.GetExperiment()
        parse_dict(input_dict, request_message)
        return mlflow.server.handlers.get_experiment_impl(request_message).experiment

    def resolve_model_versions(self, info):
        run_id = self.info.run_id
        input_dict = {"filter": f"run_id='{run_id}'"}
        request_message = mlflow.protos.model_registry_pb2.SearchModelVersions()
        parse_dict(input_dict, request_message)
        return mlflow.server.handlers.search_model_versions_impl(request_message).model_versions


class MlflowMetricExtension(MlflowMetric):
    value = graphene.Float()

    # metric values that are NaN will cause an error in graphQL validation as
    # the type is Float. as a workaround, we return None if the value is NaN.
    def resolve_value(self, info):
        return None if math.isnan(self.value) else self.value


class Query(QueryType):
    test = graphene.Field(Test, input_string=graphene.String(), description="Simple echoing field")
    mlflow_search_runs = graphene.Field(MlflowSearchRunsResponse, input=MlflowSearchRunsInput())

    def resolve_test(self, info, input_string):
        return {"output": input_string}

    def resolve_mlflow_search_runs(self, info, input):
        input_dict = vars(input)
        request_message = mlflow.protos.service_pb2.SearchRuns()
        parse_dict(input_dict, request_message)
        return mlflow.server.handlers.search_runs_impl(request_message)


class Mutation(MutationType):
    testMutation = graphene.Field(
        TestMutation, input_string=graphene.String(), description="Simple echoing field"
    )

    def resolve_test_mutation(self, info, input_string):
        return {"output": input_string}


schema = graphene.Schema(query=Query, mutation=Mutation, directives=[ComponentDirective])