# Copyright 2014 Cloudera Inc.
#
# 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.

from typing import Optional

import pytest

import ibis.expr.types as ir
from ibis.backends.base.sql import BaseSQLBackend
from ibis.backends.base.sql.alchemy import (
    AlchemyCompiler,
    AlchemyTable,
    table_from_schema,
)
from ibis.expr.schema import Schema
from ibis.expr.typing import TimeContext

MOCK_TABLES = {
    'alltypes': [
        ('a', 'int8'),
        ('b', 'int16'),
        ('c', 'int32'),
        ('d', 'int64'),
        ('e', 'float'),
        ('f', 'double'),
        ('g', 'string'),
        ('h', 'boolean'),
        ('i', 'timestamp'),
        ('j', 'date'),
        ('k', 'time'),
    ],
    'star1': [
        ('c', 'int32'),
        ('f', 'double'),
        ('foo_id', 'string'),
        ('bar_id', 'string'),
    ],
    'star2': [
        ('foo_id', 'string'),
        ('value1', 'double'),
        ('value3', 'double'),
    ],
    'star3': [('bar_id', 'string'), ('value2', 'double')],
    'test1': [('c', 'int32'), ('f', 'double'), ('g', 'string')],
    'test2': [('key', 'string'), ('value', 'double')],
    'tpch_region': [
        ('r_regionkey', 'int16'),
        ('r_name', 'string'),
        ('r_comment', 'string'),
    ],
    'tpch_nation': [
        ('n_nationkey', 'int16'),
        ('n_name', 'string'),
        ('n_regionkey', 'int16'),
        ('n_comment', 'string'),
    ],
    'tpch_lineitem': [
        ('l_orderkey', 'int64'),
        ('l_partkey', 'int64'),
        ('l_suppkey', 'int64'),
        ('l_linenumber', 'int32'),
        ('l_quantity', 'decimal(12,2)'),
        ('l_extendedprice', 'decimal(12,2)'),
        ('l_discount', 'decimal(12,2)'),
        ('l_tax', 'decimal(12,2)'),
        ('l_returnflag', 'string'),
        ('l_linestatus', 'string'),
        ('l_shipdate', 'string'),
        ('l_commitdate', 'string'),
        ('l_receiptdate', 'string'),
        ('l_shipinstruct', 'string'),
        ('l_shipmode', 'string'),
        ('l_comment', 'string'),
    ],
    'tpch_customer': [
        ('c_custkey', 'int64'),
        ('c_name', 'string'),
        ('c_address', 'string'),
        ('c_nationkey', 'int16'),
        ('c_phone', 'string'),
        ('c_acctbal', 'decimal'),
        ('c_mktsegment', 'string'),
        ('c_comment', 'string'),
    ],
    'tpch_orders': [
        ('o_orderkey', 'int64'),
        ('o_custkey', 'int64'),
        ('o_orderstatus', 'string'),
        ('o_totalprice', 'decimal(12,2)'),
        ('o_orderdate', 'string'),
        ('o_orderpriority', 'string'),
        ('o_clerk', 'string'),
        ('o_shippriority', 'int32'),
        ('o_comment', 'string'),
    ],
    'functional_alltypes': [
        ('id', 'int32'),
        ('bool_col', 'boolean'),
        ('tinyint_col', 'int8'),
        ('smallint_col', 'int16'),
        ('int_col', 'int32'),
        ('bigint_col', 'int64'),
        ('float_col', 'float'),
        ('double_col', 'double'),
        ('date_string_col', 'string'),
        ('string_col', 'string'),
        ('timestamp_col', 'timestamp'),
        ('year', 'int32'),
        ('month', 'int32'),
    ],
    'airlines': [
        ('year', 'int32'),
        ('month', 'int32'),
        ('day', 'int32'),
        ('dayofweek', 'int32'),
        ('dep_time', 'int32'),
        ('crs_dep_time', 'int32'),
        ('arr_time', 'int32'),
        ('crs_arr_time', 'int32'),
        ('carrier', 'string'),
        ('flight_num', 'int32'),
        ('tail_num', 'int32'),
        ('actual_elapsed_time', 'int32'),
        ('crs_elapsed_time', 'int32'),
        ('airtime', 'int32'),
        ('arrdelay', 'int32'),
        ('depdelay', 'int32'),
        ('origin', 'string'),
        ('dest', 'string'),
        ('distance', 'int32'),
        ('taxi_in', 'int32'),
        ('taxi_out', 'int32'),
        ('cancelled', 'int32'),
        ('cancellation_code', 'string'),
        ('diverted', 'int32'),
        ('carrier_delay', 'int32'),
        ('weather_delay', 'int32'),
        ('nas_delay', 'int32'),
        ('security_delay', 'int32'),
        ('late_aircraft_delay', 'int32'),
    ],
    'tpcds_customer': [
        ('c_customer_sk', 'int64'),
        ('c_customer_id', 'string'),
        ('c_current_cdemo_sk', 'int32'),
        ('c_current_hdemo_sk', 'int32'),
        ('c_current_addr_sk', 'int32'),
        ('c_first_shipto_date_sk', 'int32'),
        ('c_first_sales_date_sk', 'int32'),
        ('c_salutation', 'string'),
        ('c_first_name', 'string'),
        ('c_last_name', 'string'),
        ('c_preferred_cust_flag', 'string'),
        ('c_birth_day', 'int32'),
        ('c_birth_month', 'int32'),
        ('c_birth_year', 'int32'),
        ('c_birth_country', 'string'),
        ('c_login', 'string'),
        ('c_email_address', 'string'),
        ('c_last_review_date', 'string'),
    ],
    'tpcds_customer_address': [
        ('ca_address_sk', 'bigint'),
        ('ca_address_id', 'string'),
        ('ca_street_number', 'string'),
        ('ca_street_name', 'string'),
        ('ca_street_type', 'string'),
        ('ca_suite_number', 'string'),
        ('ca_city', 'string'),
        ('ca_county', 'string'),
        ('ca_state', 'string'),
        ('ca_zip', 'string'),
        ('ca_country', 'string'),
        ('ca_gmt_offset', 'decimal(5,2)'),
        ('ca_location_type', 'string'),
    ],
    'tpcds_customer_demographics': [
        ('cd_demo_sk', 'bigint'),
        ('cd_gender', 'string'),
        ('cd_marital_status', 'string'),
        ('cd_education_status', 'string'),
        ('cd_purchase_estimate', 'int'),
        ('cd_credit_rating', 'string'),
        ('cd_dep_count', 'int'),
        ('cd_dep_employed_count', 'int'),
        ('cd_dep_college_count', 'int'),
    ],
    'tpcds_date_dim': [
        ('d_date_sk', 'bigint'),
        ('d_date_id', 'string'),
        ('d_date', 'string'),
        ('d_month_seq', 'int'),
        ('d_week_seq', 'int'),
        ('d_quarter_seq', 'int'),
        ('d_year', 'int'),
        ('d_dow', 'int'),
        ('d_moy', 'int'),
        ('d_dom', 'int'),
        ('d_qoy', 'int'),
        ('d_fy_year', 'int'),
        ('d_fy_quarter_seq', 'int'),
        ('d_fy_week_seq', 'int'),
        ('d_day_name', 'string'),
        ('d_quarter_name', 'string'),
        ('d_holiday', 'string'),
        ('d_weekend', 'string'),
        ('d_following_holiday', 'string'),
        ('d_first_dom', 'int'),
        ('d_last_dom', 'int'),
        ('d_same_day_ly', 'int'),
        ('d_same_day_lq', 'int'),
        ('d_current_day', 'string'),
        ('d_current_week', 'string'),
        ('d_current_month', 'string'),
        ('d_current_quarter', 'string'),
        ('d_current_year', 'string'),
    ],
    'tpcds_household_demographics': [
        ('hd_demo_sk', 'bigint'),
        ('hd_income_band_sk', 'int'),
        ('hd_buy_potential', 'string'),
        ('hd_dep_count', 'int'),
        ('hd_vehicle_count', 'int'),
    ],
    'tpcds_item': [
        ('i_item_sk', 'bigint'),
        ('i_item_id', 'string'),
        ('i_rec_start_date', 'string'),
        ('i_rec_end_date', 'string'),
        ('i_item_desc', 'string'),
        ('i_current_price', 'decimal(7,2)'),
        ('i_wholesale_cost', 'decimal(7,2)'),
        ('i_brand_id', 'int'),
        ('i_brand', 'string'),
        ('i_class_id', 'int'),
        ('i_class', 'string'),
        ('i_category_id', 'int'),
        ('i_category', 'string'),
        ('i_manufact_id', 'int'),
        ('i_manufact', 'string'),
        ('i_size', 'string'),
        ('i_formulation', 'string'),
        ('i_color', 'string'),
        ('i_units', 'string'),
        ('i_container', 'string'),
        ('i_manager_id', 'int'),
        ('i_product_name', 'string'),
    ],
    'tpcds_promotion': [
        ('p_promo_sk', 'bigint'),
        ('p_promo_id', 'string'),
        ('p_start_date_sk', 'int'),
        ('p_end_date_sk', 'int'),
        ('p_item_sk', 'int'),
        ('p_cost', 'decimal(15,2)'),
        ('p_response_target', 'int'),
        ('p_promo_name', 'string'),
        ('p_channel_dmail', 'string'),
        ('p_channel_email', 'string'),
        ('p_channel_catalog', 'string'),
        ('p_channel_tv', 'string'),
        ('p_channel_radio', 'string'),
        ('p_channel_press', 'string'),
        ('p_channel_event', 'string'),
        ('p_channel_demo', 'string'),
        ('p_channel_details', 'string'),
        ('p_purpose', 'string'),
        ('p_discount_active', 'string'),
    ],
    'tpcds_store': [
        ('s_store_sk', 'bigint'),
        ('s_store_id', 'string'),
        ('s_rec_start_date', 'string'),
        ('s_rec_end_date', 'string'),
        ('s_closed_date_sk', 'int'),
        ('s_store_name', 'string'),
        ('s_number_employees', 'int'),
        ('s_floor_space', 'int'),
        ('s_hours', 'string'),
        ('s_manager', 'string'),
        ('s_market_id', 'int'),
        ('s_geography_class', 'string'),
        ('s_market_desc', 'string'),
        ('s_market_manager', 'string'),
        ('s_division_id', 'int'),
        ('s_division_name', 'string'),
        ('s_company_id', 'int'),
        ('s_company_name', 'string'),
        ('s_street_number', 'string'),
        ('s_street_name', 'string'),
        ('s_street_type', 'string'),
        ('s_suite_number', 'string'),
        ('s_city', 'string'),
        ('s_county', 'string'),
        ('s_state', 'string'),
        ('s_zip', 'string'),
        ('s_country', 'string'),
        ('s_gmt_offset', 'decimal(5,2)'),
        ('s_tax_precentage', 'decimal(5,2)'),
    ],
    'tpcds_store_sales': [
        ('ss_sold_time_sk', 'bigint'),
        ('ss_item_sk', 'bigint'),
        ('ss_customer_sk', 'bigint'),
        ('ss_cdemo_sk', 'bigint'),
        ('ss_hdemo_sk', 'bigint'),
        ('ss_addr_sk', 'bigint'),
        ('ss_store_sk', 'bigint'),
        ('ss_promo_sk', 'bigint'),
        ('ss_ticket_number', 'int'),
        ('ss_quantity', 'int'),
        ('ss_wholesale_cost', 'decimal(7,2)'),
        ('ss_list_price', 'decimal(7,2)'),
        ('ss_sales_price', 'decimal(7,2)'),
        ('ss_ext_discount_amt', 'decimal(7,2)'),
        ('ss_ext_sales_price', 'decimal(7,2)'),
        ('ss_ext_wholesale_cost', 'decimal(7,2)'),
        ('ss_ext_list_price', 'decimal(7,2)'),
        ('ss_ext_tax', 'decimal(7,2)'),
        ('ss_coupon_amt', 'decimal(7,2)'),
        ('ss_net_paid', 'decimal(7,2)'),
        ('ss_net_paid_inc_tax', 'decimal(7,2)'),
        ('ss_net_profit', 'decimal(7,2)'),
        ('ss_sold_date_sk', 'bigint'),
    ],
    'tpcds_time_dim': [
        ('t_time_sk', 'bigint'),
        ('t_time_id', 'string'),
        ('t_time', 'int'),
        ('t_hour', 'int'),
        ('t_minute', 'int'),
        ('t_second', 'int'),
        ('t_am_pm', 'string'),
        ('t_shift', 'string'),
        ('t_sub_shift', 'string'),
        ('t_meal_time', 'string'),
    ],
}


class MockBackend(BaseSQLBackend):
    name = 'mock'
    version = '1.0'
    current_database = 'mockdb'

    def __init__(self):
        self.executed_queries = []

    def connect(self):
        pass

    def list_tables(self):
        return list(MOCK_TABLES)

    def list_databases(self):
        return ['mockdb']

    def fetch_from_cursor(self, cursor, schema):
        pass

    def get_schema(self, name):
        name = name.replace('`', '')
        return Schema.from_tuples(MOCK_TABLES[name])

    def execute(self, expr, limit=None, params=None, **kwargs):
        ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
        for query in ast.queries:
            self.executed_queries.append(query.compile())
        return None

    def compile(
        self,
        expr,
        limit=None,
        params=None,
        timecontext: Optional[TimeContext] = None,
    ):
        ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
        queries = [q.compile() for q in ast.queries]
        return queries[0] if len(queries) == 1 else queries


class MockAlchemyBackend(MockBackend):
    compiler = AlchemyCompiler

    def __init__(self):
        super().__init__()
        sa = pytest.importorskip('sqlalchemy')
        self.meta = sa.MetaData()

    def table(self, name, database=None):
        schema = self.get_schema(name)
        return self._inject_table(name, schema)

    def _inject_table(self, name, schema):
        if name in self.meta.tables:
            table = self.meta.tables[name]
        else:
            table = table_from_schema(name, self.meta, schema)

        node = AlchemyTable(table, self)
        return ir.TableExpr(node)


GEO_TABLE = {
    'geo': [
        ('id', 'int32'),
        ('geo_point', 'point'),
        ('geo_linestring', 'linestring'),
        ('geo_polygon', 'polygon'),
        ('geo_multipolygon', 'multipolygon'),
    ]
}


class GeoMockConnectionPostGIS(MockAlchemyBackend):
    _tables = GEO_TABLE

    def __init__(self):
        super().__init__()
        self.executed_queries = []

    def get_schema(self, name):
        return Schema.from_tuples(self._tables[name])


class GeoMockConnectionOmniSciDB(MockBackend):
    _tables = GEO_TABLE

    def __init__(self):
        super().__init__()
        self.executed_queries = []

    def get_schema(self, name):
        return Schema.from_tuples(self._tables[name])
