import os

from autogeneration_utils import AUTOGENERATED_SCHEMA, AUTOGENERATED_SDL_SCHEMA
from parsing_utils import process_method
from schema_autogeneration import generate_schema

from mlflow.protos import model_registry_pb2, service_pb2
from mlflow.server.graphql.graphql_schema_extensions import schema

# Add proto descriptors to onboard RPCs to graphql.
ONBOARDED_DESCRIPTORS = [service_pb2.DESCRIPTOR, model_registry_pb2.DESCRIPTOR]


class GenerateSchemaState:
    def __init__(self):
        self.queries = set()  # method_descriptor
        self.mutations = set()  # method_descriptor
        self.inputs = []  # field_descriptor
        self.outputs = set()  # field_descriptor
        self.types = []  # field_descriptor
        self.enums = set()  # enum_descriptor
        self.method_names = set()  # package_name_method_name


# Entry point for generating the GraphQL schema.
def generate_code():
    state = GenerateSchemaState()
    for file_descriptor in ONBOARDED_DESCRIPTORS:
        for service_name, service_descriptor in file_descriptor.services_by_name.items():
            for method_name, method_descriptor in service_descriptor.methods_by_name.items():
                process_method(method_descriptor, state)

    generated_schema = generate_schema(state)

    os.makedirs(os.path.dirname(AUTOGENERATED_SCHEMA), exist_ok=True)

    with open(AUTOGENERATED_SCHEMA, "w") as file:
        file.write(generated_schema)

    # Generate the sdl schema for typescript type generation.
    sdl_schema = str(schema)
    sdl_schema = f"""# GENERATED FILE. PLEASE DON'T MODIFY.
# Run uv run ./dev/proto_to_graphql/code_generator.py to regenerate.

{sdl_schema}

"""

    with open(AUTOGENERATED_SDL_SCHEMA, "w") as f:
        f.write(sdl_schema)


def main():
    generate_code()


if __name__ == "__main__":
    main()