from collections import defaultdict
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from pprint import pprint
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Union,
)

from dagster import (
    AssetExecutionContext,
    AssetsDefinition,
    AssetSelection,
    HourlyPartitionsDefinition,
    ScheduleDefinition,
    asset,
    build_schedule_from_partitioned_job,
    define_asset_job,
    fs_io_manager,
)

from .dagster_snowflake_insights import (
    AssetMaterializationId,
    get_cost_data_for_hour,
)
from .metrics_utils import DagsterMetric, put_metrics
from .snowflake.snowflake_utils import OUTPUT_NON_ASSET_SIGIL

if TYPE_CHECKING:
    from dagster_snowflake import SnowflakeConnection

SNOWFLAKE_QUERY_HISTORY_LATENCY_SLA_MINS = 45


@dataclass
class SnowflakeInsightsDefinitions:
    assets: Sequence[AssetsDefinition]
    schedule: ScheduleDefinition


def create_snowflake_insights_asset_and_schedule(
    start_date: Union[datetime, date, str],
    name: Optional[str] = None,
    group_name: Optional[str] = None,
    job_name: str = "snowflake_insights_import",
    dry_run=False,
    allow_partial_partitions=False,
    snowflake_resource_key: str = "snowflake",
    snowflake_usage_latency: int = SNOWFLAKE_QUERY_HISTORY_LATENCY_SLA_MINS,
    partition_end_offset_hrs: int = 1,
) -> SnowflakeInsightsDefinitions:
    """Generates a pre-defined Dagster asset and schedule that can be used to import Snowflake cost
    data into Dagster Insights.

    The schedule will run hourly, and will query the Snowflake query_history table for all queries
    that ran in the hour starting at the scheduled time. It will then submit the cost data to
    Dagster Insights.

    Args:
        start_date (Union[datetime, str]): The date to start the partitioned schedule on. This should be the date
            that you began to track cost data alongside your dbt runs.
        name (Optional[str]): The name of the asset. Defaults to "snowflake_query_history".
        group_name (Optional[str]): The name of the asset group. Defaults to the default group.
        job_name (str): The name of the job that will be created to run the schedule. Defaults to
            "snowflake_insights_import".
        dry_run (bool): If true, the schedule will print the cost data to stdout instead of
            submitting it to Dagster Insights. Defaults to True.
        snowflake_resource_key (str): The name of the snowflake resource key to use. Defaults to
            "snowflake".
        partition_end_offset_hrs (int): The number of additional hours to wait before querying
            Snowflake for the latest data. This is useful in case the Snowflake query_history table
            is not immediately available. Defaults to .
    """
    # for backcompat, this used to take `date`
    if isinstance(start_date, date):
        start_date = start_date.strftime("%Y-%m-%d-%H:%M")

    @asset(
        name=name,
        group_name=group_name,
        partitions_def=HourlyPartitionsDefinition(
            start_date=start_date, end_offset=partition_end_offset_hrs
        ),
        required_resource_keys={snowflake_resource_key},
        io_manager_def=fs_io_manager,
    )
    def poll_snowflake_query_history_hour(
        context: AssetExecutionContext,
    ) -> None:
        snowflake: "SnowflakeConnection" = getattr(context.resources, snowflake_resource_key)

        start_hour = context.partition_time_window.start
        end_hour = context.partition_time_window.end

        if not allow_partial_partitions:
            now = datetime.now().astimezone(timezone.utc)
            earliest_call_time = end_hour + timedelta(minutes=snowflake_usage_latency)
            if now < earliest_call_time:
                raise RuntimeError(
                    "This function was called before the Snowflake query_history table may be"
                    f" available. For hour starting {start_hour.isoformat()} you can call it"
                    f" starting at {earliest_call_time.isoformat()} (it is currently"
                    f" {now.isoformat()})"
                )

        costs = (
            get_cost_data_for_hour(
                snowflake,
                context.instance,
                start_hour,
                end_hour,
            )
            or []
        )
        if dry_run:
            pprint(costs)
        else:
            mat_by_step: Mapping[
                Tuple[str, str], List[Tuple[AssetMaterializationId, Any]]
            ] = defaultdict(list)
            for mat, cost in costs:
                mat_by_step[(mat.run_id, mat.step_key)].append((mat, cost))

            for run_id, step_key in mat_by_step.keys():
                materializations = mat_by_step[(run_id, step_key)]
                costs_by_asset_key_and_partition: Dict[Tuple[str, Optional[str]], float] = {}
                for mat, cost in materializations:
                    # filter out any costs not associated with a specific asset (e.g. just an op output)
                    if not mat.asset_key.path[-1].startswith(OUTPUT_NON_ASSET_SIGIL):
                        key = (mat.asset_key.to_string(), mat.partition)
                        costs_by_asset_key_and_partition[key] = (
                            costs_by_asset_key_and_partition.get(key, 0) + float(cost)
                        )

                costs_keys = list(costs_by_asset_key_and_partition.keys())

                for asset_key, partition in costs_keys:
                    cost = costs_by_asset_key_and_partition[(asset_key, partition)]

                    result = put_metrics(
                        context=context,
                        run_id=run_id,
                        step_key=step_key,
                        asset_key=asset_key,
                        partition=partition,
                        metrics=[DagsterMetric("snowflake_credits", cost)],
                    )
                    assert (
                        result["createOrUpdateMetrics"]["__typename"]
                        == "CreateOrUpdateMetricsSuccess"
                    ), result

                result = put_metrics(
                    context=context,
                    run_id=run_id,
                    step_key=step_key,
                    metrics=[
                        DagsterMetric(
                            "snowflake_credits",
                            sum([float(cost) for mat, cost in materializations]),
                        )
                    ],
                )
                assert (
                    result["createOrUpdateMetrics"]["__typename"] == "CreateOrUpdateMetricsSuccess"
                ), result

    schedule = build_schedule_from_partitioned_job(
        job=define_asset_job(job_name, AssetSelection.assets(poll_snowflake_query_history_hour)),
        minute_of_hour=59,
    )
    # schedule may be a UnresolvedPartitionedAssetScheduleDefinition so we ignore the type check
    return SnowflakeInsightsDefinitions(
        assets=[poll_snowflake_query_history_hour],
        schedule=schedule,  # type: ignore
    )
