import functools
from typing import Optional, Any, Dict, Union

from dagshub.data_engine.client.query_builder import GqlQuery


class GqlQueries:
    @staticmethod
    @functools.lru_cache()
    def datasource() -> str:
        q = (
            GqlQuery()
            .operation("query", name="datasource", input={"$id": "ID", "$name": "String"})
            .query(
                "datasource",
                input={
                    "id": "$id",
                    "name": "$name",
                },
            )
            .fields(
                [
                    "id",
                    "name",
                    "rootUrl",
                    "integrationStatus",
                    "preprocessingStatus",
                    "metadataFields {name valueType multiple tags}" "type",
                ]
            )
            .generate()
        )
        return q

    @staticmethod
    def datasource_params(id: Optional[Union[int, str]], name: Optional[str]) -> Dict[str, Any]:
        return {
            "id": id,
            "name": name,
        }

    @staticmethod
    @functools.lru_cache()
    def datasource_query(include_metadata: bool) -> str:
        metadata_fields = "metadata { key value }" if include_metadata else ""
        q = (
            GqlQuery()
            .operation(
                "query",
                name="datasourceQuery",
                input={
                    "$datasource": "ID!",
                    "$queryInput": "QueryInput",
                    "$first": "Int",
                    "$after": "String",
                },
            )
            .query(
                "datasourceQuery",
                input={
                    "datasource": "$datasource",
                    "filter": "$queryInput",
                    "first": "$first",
                    "after": "$after",
                },
            )
            .fields(
                [
                    f"edges {{ node {{ id path {metadata_fields} }} }}",
                    "pageInfo { hasNextPage endCursor }",
                ]
            )
            .generate()
        )
        return q

    @staticmethod
    def datasource_query_params(
        datasource_id: Union[int, str], query_input: Dict[str, Any], first: Optional[int], after: Optional[str]
    ) -> Dict[str, Any]:
        return {
            "datasource": datasource_id,
            "queryInput": query_input,
            "first": first,
            "after": after,
        }

    @staticmethod
    @functools.lru_cache()
    def dataset() -> str:
        q = (
            GqlQuery()
            .operation("query", name="dataset", input={"$id": "ID", "$name": "String"})
            .query(
                "dataset",
                input={
                    "dataset": "$id",
                    "name": "$name",
                },
            )
            .fields(
                [
                    "id",
                    "name",
                    "datasource {id name rootUrl integrationStatus preprocessingStatus "
                    "metadataFields {name valueType multiple tags} type}",
                    "datasetQuery",
                ]
            )
            .generate()
        )
        return q

    @staticmethod
    def dataset_params(id: Optional[Union[int, str]], name: Optional[str]) -> Dict[str, Any]:
        return {
            "id": id,
            "name": name,
        }
