import re
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict
from typing import List
from typing import Optional

import pandas
import pendulum

from tecton_proto.common import aggregation_function_pb2 as afpb
from tecton_proto.data.feature_service_pb2 import FeatureSetItem
from tecton_proto.data.feature_view_pb2 import FeatureView as FeatureViewProto
from tecton_proto.data.feature_view_pb2 import NewTemporalAggregate
from tecton_spark.snowflake_pipeline_helper import pipeline_to_sql_string
from tecton_spark.templates_utils import load_template

HISTORICAL_FEATURES_TEMPLATE = None
MATERIALIZATION_TILE_TEMPLATE = None
MATERIALIZATION_TEMPLATE = None

# TODO(TEC-6204): Last and LastN are not currently supported
AGGREGATION_PLANS = {
    afpb.AGGREGATION_FUNCTION_SUM: {"SUM"},
    afpb.AGGREGATION_FUNCTION_MIN: {"MIN"},
    afpb.AGGREGATION_FUNCTION_MAX: {"MAX"},
    afpb.AGGREGATION_FUNCTION_COUNT: {"COUNT"},
    afpb.AGGREGATION_FUNCTION_MEAN: {"COUNT", "SUM"},
}


def _load_template():
    # TODO: Do this at module loading time once we sort out including the templates in the public SDK build
    global HISTORICAL_FEATURES_TEMPLATE
    if not HISTORICAL_FEATURES_TEMPLATE:
        HISTORICAL_FEATURES_TEMPLATE = load_template("historical_features.sql.j2")
    global MATERIALIZATION_TILE_TEMPLATE
    if not MATERIALIZATION_TILE_TEMPLATE:
        MATERIALIZATION_TILE_TEMPLATE = load_template("materialization_tile.sql.j2")
    global MATERIALIZATION_TEMPLATE
    if not MATERIALIZATION_TEMPLATE:
        MATERIALIZATION_TEMPLATE = load_template("materialization.sql.j2")


@dataclass
class _FeatureSetItemInput:
    """A simplified version of FeatureSetItem which is passed to the SQL template."""

    name: str
    timestamp_key: str
    join_keys: Dict[str, str]
    features: List[str]
    sql: str
    aggregation: Optional[NewTemporalAggregate]


def get_historical_features(
    spine: pandas.DataFrame,
    feature_set_items: List[FeatureSetItem],
    timestamp_key: str,
    username: str,
    password: str,
    include_feature_package_timestamp_columns: bool = False,
) -> pandas.DataFrame:
    # snowflake-connector-python is not currently required as a dependency due to pyarrow conflict with pyspark
    # To compile, see a working setup here: https://tecton-ai.phacility.com/D9593?id=32124
    import snowflake.connector
    from tecton_spark.snowflake_pandas_tools import write_pandas

    # TODO: Currently we try to fetch the first batch data source for the first feature view,
    # may need to think about a better way to get the warehouse/database/schema.
    data_sources = list(feature_set_items[0].enrichments.feature_view.enrichments.virtual_data_sources)
    for ds in data_sources:
        if ds.HasField("batch_data_source") and ds.batch_data_source.HasField("snowflake"):
            snowflake_ds = ds.batch_data_source.snowflake
            snowflake_warehouse = snowflake_ds.snowflakeArgs.warehouse
            snowflake_database = snowflake_ds.snowflakeArgs.database
            snowflake_schema = snowflake_ds.snowflakeArgs.schema

    # Warehouse, database, schema is needed to store temporary table.
    # TODO: Currently we hard coded the account part here, need to allow clients to change it.
    conn = snowflake.connector.connect(
        user=username,
        password=password,
        account="tectonpartner",
        warehouse=snowflake_warehouse,
        database=snowflake_database,
        schema=snowflake_schema,
    )
    try:
        cur = conn.cursor()
        temp_table_name = "_temp_spine_table_from_df"
        # Get the SQL that would be generated by the create table statement
        create_table_sql = pandas.io.sql.get_schema(spine, temp_table_name)

        # Replace the `CREATE TABLE` with `CREATE OR REPLACE TEMPORARY TABLE`
        create_tmp_table_sql = re.sub("^(CREATE TABLE)?", "CREATE OR REPLACE TEMPORARY TABLE", create_table_sql)

        get_historical_features_sql = get_features_sql_str_for_spine(
            feature_set_items=feature_set_items,
            timestamp_key=timestamp_key,
            # Connector will add a double quote to the table name, we have to mimic that here otherwise it will fail to
            # find the table.
            spine_sql=f'SELECT * FROM "{temp_table_name}"',
            include_feature_package_timestamp_columns=include_feature_package_timestamp_columns,
        )
        cur.execute(create_tmp_table_sql)
        # write_pandas sometimes does not parse the pandas timestamp value correctly.
        # If we create a timestamp with datetime.datetime(1997, 6, 3, 14, 21, 32, 00),
        # it will turns into this in snowflake:
        # 1997-06-03 14:21:32 -> 29391-10-10 22:53:20.000
        # This will work correctly with snowflake:
        # datetime.datetime(1997, 6, 3, 14, 21, 32, 00, tzinfo=datetime.timezone.utc)
        # I *think* what it needs to be parsed correctly is timezone, need confirmation from snowflake.
        # Related open issue: https://github.com/snowflakedb/snowflake-connector-python/issues/600
        #
        # Note: Snowflake provided a hacky fix for this(our own write_pandas function), but it may break at some point
        # because it's using a deprecated feature.
        # TODO: Using snowflake.connector.pandas_tools.write_pandas instead when it's fixed.
        # TODO: Snowflake also does not support nano seconds due to how it dump the data into database, we
        # need to see if we ever run into any issue with this and come up with a fix for it.
        write_pandas(conn=conn, df=spine, table_name=temp_table_name)
        cur.execute(get_historical_features_sql)
        df = cur.fetch_pandas_all()
    finally:
        # Closing the connection
        conn.close()
    return df


def get_materialization_sql_str(
    feature_view: FeatureViewProto,
    # start is inclusive and end is exclusive
    time_limits: pendulum.Period,
    destination: str,
    storage_integration: Optional[str],
) -> str:
    _load_template()
    source = feature_view.fco_metadata.name.replace("-", "_")
    if feature_view.HasField("temporal_aggregate"):
        aggregations = defaultdict(set)
        for feature in feature_view.temporal_aggregate.features:
            aggregate_function = AGGREGATION_PLANS[feature.function]
            if not aggregate_function:
                raise ValueError(f"Unsupported aggregation function {feature.function} in snowflake pipeline")
            aggregations[feature.input_feature_name].update(aggregate_function)

        # Need to order the functions for determinism results.
        for key, value in aggregations.items():
            aggregations[key] = sorted(value)

        source = MATERIALIZATION_TILE_TEMPLATE.render(
            source=source,
            join_keys=list(feature_view.join_keys),
            aggregations=aggregations,
            slide_interval=feature_view.temporal_aggregate.slide_interval,
            timestamp_key=feature_view.timestamp_key,
        )
    return MATERIALIZATION_TEMPLATE.render(
        source=source,
        materialization_schema=feature_view.schemas.materialization_schema,
        timestamp_key=feature_view.timestamp_key,
        start_time=time_limits.start,
        end_time=time_limits.end,
        destination=destination,
        storage_integration=storage_integration,
    )


def get_features_sql_str_for_spine(
    feature_set_items: List[FeatureSetItem],
    timestamp_key: str,
    spine_sql: str = None,
    include_feature_package_timestamp_columns: bool = False,
) -> str:
    """
    Get a SQL string to fetch features given the spine and feature set.
    spine_sql and spine_table_name cannot be both empty.

    :param feature_set_items: FeatureSetItems for the features.
    :param timestamp_key: Name of the time column in the spine.
    :param spine_sql: SQL str to get the spine.
    :param include_feature_package_timestamp_columns: (Optional) Include timestamp columns for every individual FeaturePackage.
    :return: A SQL string that can be used to fetch features.
    """
    _load_template()
    if include_feature_package_timestamp_columns:
        raise NotImplementedError()

    input_items = []
    for item in feature_set_items:
        feature_view = item.enrichments.feature_view
        join_keys = {k.package_column_name: k.spine_column_name for k in item.join_configuration_items}
        features = [
            col.name
            for col in feature_view.schemas.view_schema.columns
            if col.name not in (list(join_keys.keys()) + [feature_view.timestamp_key])
        ]
        if len(feature_view.online_serving_index.join_keys) != len(feature_view.join_keys):
            raise ValueError("SQL string does not support wildcard")
        input_items.append(
            _FeatureSetItemInput(
                name=feature_view.fco_metadata.name,
                timestamp_key=feature_view.timestamp_key,
                join_keys=join_keys,
                features=features,
                sql=pipeline_to_sql_string(
                    pipeline=feature_view.pipeline,
                    virtual_data_sources=feature_view.enrichments.virtual_data_sources,
                    transformations=feature_view.enrichments.transformations,
                ),
                aggregation=(
                    feature_view.temporal_aggregate
                    if item.enrichments.feature_view.HasField("temporal_aggregate")
                    else None
                ),
            )
        )
    return HISTORICAL_FEATURES_TEMPLATE.render(
        feature_set_items=input_items,
        spine_timestamp_key=timestamp_key,
        spine_sql=spine_sql,
        include_feature_package_timestamp_columns=include_feature_package_timestamp_columns,
    )
