# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ibis
from data_validation import consts

from datetime import datetime, timedelta


class AggregateField(object):
    def __init__(self, ibis_expr, field_name=None, name=None):
        """
            field_name: A field to act on in the table.  Table level expr do not have a field name
        """
        self.expr = ibis_expr
        self.field_name = field_name
        self.name = name

    @staticmethod
    def count(name=None):
        return AggregateField(
            ibis.expr.types.ColumnExpr.count, field_name=None, name=name
        )

    def compile(self, ibis_table):
        if self.field_name:
            agg_field = self.expr(ibis_table[self.field_name])
        else:
            agg_field = self.expr(ibis_table)

        if self.name:
            agg_field = agg_field.name(self.name)

        return agg_field


class FilterField(object):
    def __init__(
        self, ibis_expr, left=None, right=None, left_field=None, right_field=None
    ):
        """
            field_name: A field to act on in the table.  Table level expr do not have a field name
        """
        self.expr = ibis_expr
        self.left = left
        self.right = right
        self.left_field = left_field
        self.right_field = right_field

    @staticmethod
    def greater_than(field_name, value):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.__gt__, left_field=field_name, right=value
        )

    @staticmethod
    def less_than(field_name, value):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.__lt__, left_field=field_name, right=value
        )

    def compile(self, ibis_table):
        if self.left_field:
            self.left = ibis_table[self.left_field]
            # Cast All Datetime to Date (TODO this may be a bug in BQ)
            if isinstance(
                ibis_table[self.left_field].type(), ibis.expr.datatypes.Timestamp
            ):
                self.left = self.left.cast("date")
        if self.right_field:
            self.right = ibis_table[self.right_field]
            # Cast All Datetime to Date (TODO this may be a bug in BQ)
            if isinstance(
                ibis_table[self.right_field].type(), ibis.expr.datatypes.Timestamp
            ):
                self.right = self.right.cast("date")

        return self.expr(self.left, self.right)


class GroupedField(object):
    def __init__(self, field_name, alias=None, cast=None):
        """
            :param field_name: A field to act on in the table
            :param alias: An alias to use for the group
            :param cast: A cast on the column if required
        """
        self.field_name = field_name
        self.alias = alias
        self.cast = cast

    def compile(self, ibis_table):
        # Fields are supplied on compile or on build
        group_field = ibis_table[self.field_name]

        # TODO: generate cast for known types not specified
        if self.cast:
            group_field = group_field.cast(self.cast)
        elif isinstance(group_field.type(), ibis.expr.datatypes.Timestamp):
            group_field = group_field.cast("date")
        else:
            # TODO: need to build Truncation Int support
            # TODO: should be using a logger
            print("WARNING: Unknown cast types can cause memory errors")

        # The Casts require we also supply a name.
        alias = self.alias or self.field_name
        group_field = group_field.name(alias)

        return group_field


class QueryBuilder(object):
    def __init__(
        self, aggregate_fields, filters, grouped_fields, days_past=None, limit=None
    ):
        """ Build a QueryBuilder object which can be used to build queries easily

            :param aggregate_fields: List of AggregateField objects with Ibis expressions
            :param filters: A list of FilterField objects
            :param grouped_fields: A list of GroupedField objects
            :param days_past: Number of days past to use as date partition filter
            :param limit: A limit value for the number of records to pull (used for testing)
        """
        self.aggregate_fields = aggregate_fields
        self.filters = filters
        self.grouped_fields = grouped_fields
        self.days_past = days_past
        self.limit = limit

    @staticmethod
    def build_count_validator(limit=None):
        """ Return a basic template builder for most validations """
        aggregate_fields = [AggregateField.count("count")]
        filters = []
        grouped_fields = []

        return QueryBuilder(
            aggregate_fields,
            filters=filters,
            grouped_fields=grouped_fields,
            limit=limit,
        )

    def compile_aggregate_fields(self, table):
        aggs = [field.compile(table) for field in self.aggregate_fields]

        return aggs

    def compile_filter_fields(self, table):
        return [field.compile(table) for field in self.filters]

    def compile_group_fields(self, table):
        return [
            field.compile(table)
            for field in self.grouped_fields
        ]

    def compile(self, data_client, schema_name, table_name):
        """ Return an Ibis query object

            :IbisClient data_client: The client used to validate the query
            :String schema_name: The name of the schema for the givne table
            :String table_name: The name of the table to query
        """
        table = data_client.table(table_name, database=schema_name)

        # Build Query Expressions
        aggs = self.compile_aggregate_fields(table)
        filters = self.compile_filter_fields(table)
        groups = self.compile_group_fields(table)

        query = table.filter(filters)
        query = query.groupby(groups)
        query = query.aggregate(aggs)

        # if groups:
        #     query = table.groupby(groups).aggregate(aggs)
        # else:
        #     query = table.aggregate(aggs)

        if self.limit:
            query = query.limit(self.limit)

        return query

    def add_query_field(self, query_field):
        """ Add a QueryField object to the query

            :param query_field: A QueryField Object
        """
        self.query_fields.append(query_field)

    def add_grouped_field(self, grouped_field):
        """ Add a GroupedField object to the query

            :param query_field: A GroupedField Object
        """
        self.grouped_fields.append(grouped_field)

    def add_filter_field(self, filter_obj):
        """ Add a filter object to your query

            :param filter_obj: A FilterField Object
        """
        self.filters.append(filter_obj)

@@ -585,85 +585,8 @@
 bis%0A
-from data_validation import consts%0A%0Afrom datetime import datetime, timedelta%0A
 %0A%0Acl

package cleanup