
from dataclasses import dataclass
import dataclasses
import json
from typing import List, Optional
from sgqlc.operation import Operation
from schema import *
from cognite.client import CogniteClient


CDF_ALPHA_VERSION_HEADERS = {'cdf-version': 'alpha'}


def get_type_key(typeName: str, version: int):
    return f"{typeName}_{version}"


def get_relationship_key(typeName: str, attr: str, version: int):
    return f"{typeName}_{attr}_{version}"


@dataclass
class RelationshipItem:
    """Class for keeping track of an item in inventory."""
    start_node: str
    end_node: str


class FDMQueryClient:
    @classmethod
    def from_token(
            cls,
            token: Optional[str] = BEARER_TOKEN,
            tenant: Optional[str] = TENANT,
            project: Optional[str] = PROJECT,
            cluster: Optional[str] = CLUSTER,
            client_name: Optional[str] = "FDMQueryClient"
    ):
        return cls(client=CogniteClient(
            token_url=f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token",
            token=token,
            token_scopes=f"https://{cluster}.cognitedata.com/.default",
            project=project,
            base_url=f"https://{cluster}.cognitedata.com",
            client_name=client_name,
        ))

    @classmethod
    def from_client_secret(
            cls,
            client_id: str,
            client_secret: str,
            tenant: Optional[str] = TENANT,
            project: Optional[str] = PROJECT,
            cluster: Optional[str] = CLUSTER,
            client_name: Optional[str] = "FDMQueryClient"
    ):
        return cls(client=CogniteClient(
            token_url=f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token",
            token_client_id=client_id,
            token_client_secret=client_secret,
            token_scopes=f"https://{cluster}.cognitedata.com/.default",
            project=project,
            base_url=f"https://{cluster}.cognitedata.com",
            client_name=client_name,
        ))

    def __init__(
        self,
        client: CogniteClient,
    ) -> None:
        self.client = client

    def start_query(self):
        return Operation(Query)

    def run_query(self, query: Operation, data_model: Optional[str] = DM_NAME, version: Optional[int] = DM_VERSION):
        if version is None:
            versions = self.get_versions(data_model)
            if len(versions) > 0:
                version = self.get_versions(data_model)[-1].get("version")
        if version is None:
            raise Exception(
                "Must have at least 1 version published in data model")

        response = self.client.get(
            f"/api/v1/projects/{self.client.config.project}/schema/api/{data_model}/{version}/graphql",
            params={'query': repr(query)})

        data = response.json()

        errors = data.get('errors')
        if errors:
            raise BaseException(errors)

        return query + data

    def upsert_nodes(self, model_type: str, items: List[INGEST_TYPES], data_model: Optional[str] = DM_NAME, version: Optional[int] = DM_VERSION):
        if model_type not in INGEST_TYPE_NAMES:
            raise Exception(
                f"Model type must be one of {INGEST_TYPE_NAMES}")
        if version is None:
            versions = self.get_versions(data_model)
            if len(versions) > 0:
                version = self.get_versions(data_model)[-1].get("version")
        if version is None:
            raise Exception(
                "Must have at least 1 version published in data model")

        if not all(str(type(x)) == model_type for x in items):
            raise Exception(
                f"We only support loading in 1 type at a time, and type must be {model_type}")

        type_key = get_type_key(model_type, version)

        body = {
            'items': list(map(lambda x: x.__to_json_value__(), items)),
            'spaceExternalId': data_model,
            'model': [data_model, type_key],
            'overwrite': True,
        }

        response = self.client.post(
            f"/api/v1/projects/{self.client.config.project}/datamodelstorage/nodes",
            json=body, headers=CDF_ALPHA_VERSION_HEADERS)

        data = response.json()

        return data

    def upsert_relationships(self, model_type: str, attr: str, items: List[RelationshipItem], data_model: Optional[str] = DM_NAME, version: Optional[int] = DM_VERSION):
        if version is None:
            versions = self.get_versions(data_model)
            if len(versions) > 0:
                version = self.get_versions(data_model)[-1].get("version")
        if version is None:
            raise Exception(
                "Must have at least 1 version published in data model")

        relationship_key = get_relationship_key(model_type, attr, version)

        items = list(map(lambda x: {
            'startNode': x.start_node,
            'endNode': x.end_node,
            'type': relationship_key,
            'dummy': '',
            'externalId': f"{x.start_node}_{x.end_node}"
        }, items))

        body = {
            'items': items,
            'spaceExternalId': data_model,
            'model': [data_model, relationship_key],
            'overwrite': True,
            'autoCreateStartNodes': True,
            'autoCreateEndNodes': True,
        }

        response = self.client.post(
            f"/api/v1/projects/{self.client.config.project}/datamodelstorage/edges",
            json=body, headers=CDF_ALPHA_VERSION_HEADERS)

        data = response.json()

        return data

    def get_versions(self, data_model: Optional[str] = DM_NAME):
        response = self.client.get(
            f"/api/v1/projects/{self.client.config.project}/schema/graphql",
            params={'query': self._get_query_string(data_model)})
        versions = response.json().get("data").get("getApisByIds")[0].get(
            "versions")
        versions.sort(key=lambda item: item.get("version"))
        return versions

    def _get_query_string(self, data_model: Optional[str] = DM_NAME) -> str:
        return """query {{
            getApisByIds(externalIds: ["{data_model}"] ) {{
                externalId
                name
                description
                createdTime
                versions {{
                version
                createdTime
                dataModel {{
                    graphqlRepresentation
                }}
                }}
            }}
        }}
        """.format(data_model=data_model)


def camelCase(st):
    output = ''.join(x for x in st.title() if x.isalnum())
    return output[0].lower() + output[1:]
