from itertools import chain
from typing import Any, Dict

from lbox.exceptions import (
    InvalidAttributeError,
    InvalidQueryError,
    MalformedQueryException,
)

from labelbox import utils
from labelbox.orm.comparison import Comparison, LogicalExpression
from labelbox.orm.model import Entity, Field, Relationship

""" Common query creation functionality. """


def format_param_declaration(params):
    """Formats the parameters dictionary into a declaration of GraphQL
    query parameters.

    Args:
        params (dict): keys are query param names and values are
            (value, (field|relationship)) tuples.
    Return:
        str, the declaration of query parameters.
    """
    if not params:
        return ""

    def attr_type(attr):
        if isinstance(attr, Field):
            return attr.field_type.name
        else:
            return Field.Type.ID.name

    return (
        "("
        + ", ".join(
            "$%s: %s!" % (param, attr_type(attr))
            for param, (_, attr) in params.items()
        )
        + ")"
    )


def results_query_part(entity):
    """Generates the results part of the query. The results contain
    all the entity's fields as well as prefetched relationships.

    Note that this is a recursive function. If there is a cycle in the
    prefetched relationship graph, this function will recurse infinitely.

    Args:
        entity (type): The entity which needs fetching.
    """
    # Query for fields
    fields = [
        field.result_subquery
        if field.result_subquery is not None
        else field.graphql_name
        for field in entity.fields()
    ]

    # Query for cached relationships
    fields.extend(
        [
            Query(rel.graphql_name, rel.destination_type).format()[0]
            for rel in entity.relationships()
            if rel.cache
        ]
    )
    return " ".join(fields)


class Query:
    """A data structure used during the construction of a query. Supports
    subquery (also Query object) nesting for relationship."""

    def __init__(
        self, what, subquery, where=None, paginate=False, order_by=None
    ):
        """Initializer.
        Args:
            what (str): What is being queried. Typically an object type in
                singular or plural (i.e. "project" or "projects").
            subquery (Query or type): Either a Query object that is formatted
                recursively or a Entity subtype in which case the standard
                results (see `results_query_part`) are retrieved.
            where (None, Comparison or LogicalExpression): the filtering clause.
            paginate (bool): If the "%skip %first" pagination substring should
                be added to the query. Used for collection pagination in combination
                with PaginatedCollection.
            order_by (tuple): A tuple consisting of (Field, Field.Order) indicating
                how the query should sort the collection.
        """
        self.what = what
        self.subquery = subquery
        self.paginate = paginate
        self.where = where
        self.order_by = order_by

    def format_subquery(self):
        """Formats the subquery (a Query or Entity subtype)."""
        if isinstance(self.subquery, Query):
            return self.subquery.format()
        elif issubclass(self.subquery, Entity):
            return results_query_part(self.subquery), {}
        else:
            raise MalformedQueryException()

    def format_clauses(self, params):
        """Formats the where, order_by and pagination clauses.
        Args:
            params (dict): The current parameter dictionary.
        """

        def format_where(node):
            """Helper that resursively constructs a where clause from a
            LogicalExpression tree (leaf nodes are Comparisons)."""
            COMPARISON_TO_SUFFIX = {
                Comparison.Op.EQ: "",
                Comparison.Op.NE: "_not",
                Comparison.Op.LT: "_lt",
                Comparison.Op.GT: "_gt",
                Comparison.Op.LE: "_lte",
                Comparison.Op.GE: "_gte",
            }
            assert isinstance(node, (Comparison, LogicalExpression))
            if isinstance(node, Comparison):
                param_name = "param_%d" % len(params)
                params[param_name] = (node.value, node.field)
                return "{%s%s: $%s}" % (
                    node.field.graphql_name,
                    COMPARISON_TO_SUFFIX[node.op],
                    param_name,
                )
            if node.op == LogicalExpression.Op.NOT:
                return "{NOT: [%s]}" % format_where(node.first)

            return "{%s: [%s, %s]}" % (
                node.op.name.upper(),
                format_where(node.first),
                format_where(node.second),
            )

        paginate = "skip: %d first: %d" if self.paginate else ""

        where = "where: %s" % format_where(self.where) if self.where else ""

        if self.order_by:
            order_by = "orderBy: %s_%s" % (
                self.order_by[0].graphql_name,
                self.order_by[1].name.upper(),
            )
        else:
            order_by = ""

        clauses = " ".join(filter(None, (where, paginate, order_by)))
        return "(" + clauses + ")" if clauses else ""

    def format(self):
        """Formats the full query but without "query" prefix, query name
        and parameter declaration.
        Return:
            (str, dict) tuple. str is the query and dict maps parameter
            names to (value, field) tuples.
        """
        subquery, params = self.format_subquery()
        clauses = self.format_clauses(params)
        query = "%s%s{%s}" % (self.what, clauses, subquery)
        return query, params

    def format_top(self, name):
        """Formats the full query including "query" prefix, query name
        and parameter declaration. The result of this function can be
        sent to the Client object for execution.

        Args:
            name (str): Query name, without the "PyApi" suffix, it's appended
                automatically by this method.
        Return:
            (str, dict) tuple. str is the full query and dict maps parameter
                names to parameter values.
        """
        query, params = self.format()
        param_declaration = format_param_declaration(params)
        query = "query %sPyApi%s{%s}" % (name, param_declaration, query)
        return query, {param: value for param, (value, _) in params.items()}


def get_single(entity, uid):
    """Constructs the query and params dict for obtaining a single object. Either
    on ID, or without params.
    Args:
        entity (type): An Entity subtype being obtained.
        uid (str): The ID of the sought object. It can be None, which is legal for
            DB types that have a default object being returned (User and
            Organization).
    """
    type_name = entity.type_name()
    where = entity.uid == uid if uid else None
    return Query(utils.camel_case(type_name), entity, where).format_top(
        "Get" + type_name
    )


def logical_ops(where):
    """Returns a generator that yields all the logical operator
    type objects (`LogicalExpression.Op` instances) from a where
    clause.

    Args:
        where (LogicalExpression, Comparison or None): The where
            clause used for filtering in a query.
    Return:
        See above.
    """
    if isinstance(where, LogicalExpression):
        yield where.op
        for f in chain(logical_ops(where.first), logical_ops(where.second)):
            yield f


def check_where_clause(entity, where):
    """Checks the `where` clause of a query. A `where` clause is legal
    if it only refers to fields found in the entity it's defined for.
    Since only AND logical operations are supported server-side at the
    moment, logical OR and NOT are illegal.

    Args:
        entity (type): An Entity subclass type.
        where (LogicalExpression or Comparison): The `where` clause of
            query.
    Return:
        bool indicating if `where` is legal for `entity`.
    """

    def fields(where):
        """Yields all the fields in a `where` clause."""
        if isinstance(where, LogicalExpression):
            for f in chain(fields(where.first), fields(where.second)):
                yield f
        elif isinstance(where, Comparison):
            yield where.field

    # The `deleted` field is a special case, ignore it.
    where_fields = [f for f in fields(where) if f != Entity.deleted]
    invalid_fields = set(where_fields) - set(entity.fields())
    if invalid_fields:
        raise InvalidAttributeError(entity, invalid_fields)

    if len(set(where_fields)) != len(where_fields):
        raise InvalidQueryError(
            "Where clause contains multiple comparisons for "
            "the same field: %r." % where
        )

    if set(logical_ops(where)) not in (set(), {LogicalExpression.Op.AND}):
        raise InvalidQueryError(
            "Currently only AND logical ops are allowed in "
            "the where clause of a query."
        )


def check_order_by_clause(entity, order_by):
    """Checks that the `order_by` clause field is a part of `entity`.

    Args:
        entity (type): An Entity subclass type.
        order_by ((field, ordering)): The ordering tuple consisting of
            a field and sort ordering (ascending or descending).
    Return:
        bool indicating if `order_by` is legal for `entity`.
    """
    if order_by is not None:
        field, _ = order_by
        if field not in entity.fields():
            raise InvalidAttributeError(entity, field)


def get_all(entity, where):
    """Constructs a query that fetches all items of the given type. The
    resulting query is intended to be used for pagination, it contains
    two python-string int-placeholders (%d) for 'skip' and 'first'
    pagination parameters.

    Args:
        entity (type): The object type being queried.
        where (Comparison, LogicalExpression or None): The `where` clause
            for filtering.
    Return:
        (str, dict) tuple that is the query string and parameters.
    """
    check_where_clause(entity, where)
    type_name = entity.type_name()
    query = Query(utils.camel_case(type_name) + "s", entity, where, True)
    return query.format_top("Get" + type_name + "s")


def relationship(source, relationship, where, order_by):
    """Constructs a query that fetches all items from a -to-many
    relationship. To be used like:
        >>> project = ...
        >>> query_str, params = relationship(Project, "datasets", Dataset)
        >>> datasets = PaginatedCollection(
            client, query_str, params, ["project", "datasets"],
            Dataset)

    The resulting query is intended to be used for pagination, it contains
    two python-string int-placeholders (%d) for 'skip' and 'first'
    pagination parameters.

    Args:
        source (DbObject or type): If a `DbObject` then the source of the
            relationship (the query originates from that particular object).
            If `type`, then the source of the relationship is implicit, even
            without the ID. Used for expanding from Organization.
        relationship (Relationship): The relationship.
        where (Comparison, LogicalExpression or None): The `where` clause
            for filtering.
        order_by (None or (Field, Field.Order): The `order_by` clause for
            sorting results.
    Return:
        (str, dict) tuple that is the query string and parameters.
    """
    check_where_clause(relationship.destination_type, where)
    check_order_by_clause(relationship.destination_type, order_by)
    to_many = relationship.relationship_type == Relationship.Type.ToMany
    subquery = Query(
        relationship.graphql_name,
        relationship.destination_type,
        where,
        to_many,
        order_by,
    )
    query_where = (
        type(source).uid == source.uid if isinstance(source, Entity) else None
    )
    query = Query(utils.camel_case(source.type_name()), subquery, query_where)
    return query.format_top(
        "Get" + source.type_name() + utils.title_case(relationship.graphql_name)
    )


def create(entity, data):
    """Generates a query and parameters for creating a new DB object.

    Args:
        entity (type): An Entity subtype indicating which kind of
            DB object needs to be created.
        data (dict): A dict that maps Fields and Relationships to values, new
            object data.
    Return:
        (query_string, parameters)
    """
    type_name = entity.type_name()

    def format_param_value(attribute, param):
        if isinstance(attribute, Field):
            return "%s: $%s" % (attribute.graphql_name, param)
        else:
            return "%s: {connect: {id: $%s}}" % (
                utils.camel_case(attribute.graphql_name),
                param,
            )

    # Convert data to params
    params = {
        field.graphql_name: (value, field) for field, value in data.items()
    }

    query_str = """mutation Create%sPyApi%s{create%s(data: {%s}) {%s}} """ % (
        type_name,
        format_param_declaration(params),
        type_name,
        " ".join(
            format_param_value(attribute, param)
            for param, (_, attribute) in params.items()
        ),
        results_query_part(entity),
    )

    return query_str, {name: value for name, (value, _) in params.items()}


def update_relationship(a, b, relationship, update):
    """Updates the relationship in DB object `a` to connect or disconnect
    DB object `b`.

    Args:
        a (DbObject): The object being updated.
        b (DbObject): Object on the other side of the relationship.
        relationship (Relationship): The relationship from `a` to `b`.
        update (str): The type of update. Must be either `connect` or
            `disconnect`.
    Return:
        (query_string, query_parameters)
    """
    to_one_disconnect = (
        update == "disconnect"
        and relationship.relationship_type == Relationship.Type.ToOne
    )

    a_uid_param = utils.camel_case(type(a).type_name()) + "Id"

    if not to_one_disconnect:
        b_uid_param = utils.camel_case(type(b).type_name()) + "Id"
        param_declr = "($%s: ID!, $%s: ID!)" % (a_uid_param, b_uid_param)
        b_query = "{id: $%s}" % b_uid_param
    else:
        param_declr = "($%s: ID!)" % a_uid_param
        b_query = "true"

    query_str = """mutation %s%sAnd%sPyApi%s{update%s(
        where: {id: $%s} data: {%s: {%s: %s}}) {id}} """ % (
        utils.title_case(update),
        type(a).type_name(),
        type(b).type_name(),
        param_declr,
        utils.title_case(type(a).type_name()),
        a_uid_param,
        relationship.graphql_name,
        update,
        b_query,
    )

    if to_one_disconnect:
        params = {a_uid_param: a.uid}
    else:
        params = {a_uid_param: a.uid, b_uid_param: b.uid}

    return query_str, params


def update_fields(db_object, values):
    """Creates a query that updates `db_object` fields with the
    given values.

    Args:
        db_object (DbObject): The DB object being updated.
        values (dict): Maps Fields to new values. All Fields
            must be legit fields in `db_object`.
    Return:
        (query_string, query_parameters)
    """
    type_name = db_object.type_name()
    id_param = "%sId" % type_name
    values_str = " ".join(
        "%s: $%s" % (field.graphql_name, field.graphql_name)
        for field, _ in values.items()
    )
    params = {
        field.graphql_name: (value, field) for field, value in values.items()
    }
    params[id_param] = (db_object.uid, Entity.uid)

    query_str = """mutation update%sPyApi%s{update%s(
        where: {id: $%s} data: {%s}) {%s}} """ % (
        utils.title_case(type_name),
        format_param_declaration(params),
        type_name,
        id_param,
        values_str,
        results_query_part(type(db_object)),
    )

    return query_str, {name: value for name, (value, _) in params.items()}


def delete(db_object):
    """Generates a query that deletes the given `db_object` from the DB.

    Args:
        db_object (DbObject): The DB object being deleted.
    """
    id_param = "%sId" % db_object.type_name()
    query_str = """mutation delete%sPyApi%s{update%s(
        where: {id: $%s} data: {deleted: true}) {id}} """ % (
        db_object.type_name(),
        "($%s: ID!)" % id_param,
        db_object.type_name(),
        id_param,
    )

    return query_str, {id_param: db_object.uid}


def bulk_delete(db_objects, use_where_clause):
    """Generates a query that bulk-deletes the given `db_objects` from the
    DB.

    Args:
        db_objects (list): A list of DB objects of the same type.
        use_where_clause (bool): If the object IDs should be passed to the
            mutation in a `where` clause or directly as a mutation value.
    """
    type_name = db_objects[0].type_name()
    if use_where_clause:
        query_str = (
            "mutation delete%ssPyApi{delete%ss(where: {%sIds: [%s]}){id}}"
        )
    else:
        query_str = "mutation delete%ssPyApi{delete%ss(%sIds: [%s]){id}}"
    query_str = query_str % (
        utils.title_case(type_name),
        utils.title_case(type_name),
        utils.camel_case(type_name),
        ", ".join('"%s"' % db_object.uid for db_object in db_objects),
    )
    return query_str, {}


def where_as_dict(entity, node: Comparison) -> Dict[str, Any]:
    check_where_clause(entity, node)
    """ Helper that constructs a where clause from a Comparison node. """
    COMPARISON_TO_SUFFIX = {
        Comparison.Op.EQ: "",
        Comparison.Op.NE: "_not",
        Comparison.Op.LT: "_lt",
        Comparison.Op.GT: "_gt",
        Comparison.Op.LE: "_lte",
        Comparison.Op.GE: "_gte",
    }

    key = f"{node.field.graphql_name}{COMPARISON_TO_SUFFIX[node.op]}"
    return {key: node.value}


def order_by_as_string(entity, node: tuple) -> str:
    check_order_by_clause(entity, node)
    return f"{node[0].graphql_name}_{node[1].name.upper()}"
