import json
import logging
import os
import sys
from typing import Any, Optional

import boto3
import awswrangler as wr
from great_expectations.core.batch import BatchRequest
from great_expectations.profile.user_configurable_profiler import UserConfigurableProfiler

from psi_great_expectations.contextor import build_data_context_config
from psi_great_expectations import get_logger, log_it


class DataProfiler:
    def __init__(
        self,
        data_name: str,
        bucket: Optional[str] = None,
        prefix: Optional[str] = None,
        ignored_columns: list = [],
        excluded_expectations: list = [],
        logger: Optional[logging.Logger] = None,
        aws_session: Optional[Any] = None,
    ):
        self.__data_name = data_name
        self.__ignored_columns = ignored_columns
        self.__excluded_expectations = excluded_expectations
        self.__s3_arn = f"s3://{bucket}/{prefix}"
        self.__logger = logger or get_logger()
        self.__aws_session = aws_session or boto3.Session()
        self.__suite = f"{self.__data_name}-suite"
        self.__context = build_data_context_config(data_name)

    @log_it
    def _run_profiler(self):

        assert os.path.isdir(
            "ge_to_profile"
        ), "ge_to_profile directory must exist before running the profiler"
        assert os.path.isfile(f"ge_to_profile/{self.__data_name}_to_profile.json")

        batch_request = BatchRequest(
            datasource_name=self.__data_name,
            data_connector_name="default_inferred_data_connector_name",
            data_asset_name=f"{self.__data_name}_to_profile",
        )

        self.__context.create_expectation_suite(
            expectation_suite_name=self.__suite, overwrite_existing=True
        )

        validator = self.__context.get_validator(
            batch_request=batch_request, expectation_suite_name=self.__suite
        )

        # prompt user if no values for ignored_columns provided, check batch for dict/list types for auto ignore
        if not self.__ignored_columns:
            resp = input(
                "GE Profiler can not process List and Dict values\nYou have not entered any values for ignored_columns\nit is highly recommended to provide a list of columns to ignore\nthe Profiler is not perfect at catching these types without manual intervention\nCONTINUE?(y/n)"
            )
            if resp == "n":
                sys.exit("You have chosent to provide values for ignored_columns")
        df = validator.head(fetch_all=True)
        is_list = (df.applymap(type) == list).all()
        lists_list = [k for k, v in is_list.to_dict().items() if v == True]
        is_dict = (df.applymap(type) == dict).all()
        dicts_list = [k for k, v in is_dict.to_dict().items() if v == True]
        final_ignored_columns = list(set(self.__ignored_columns + lists_list + dicts_list))

        # cuts a lot of the fat out of the profiling, these values should be manually tweaked regardless
        # if you want to see all of these in the profiler results just pick a random one from https://greatexpectations.io/expectations
        if not self.__excluded_expectations:
            self.__excluded_expectations = [
                "expect_column_min_to_be_between",
                "expect_column_max_to_be_between",
                "expect_column_mean_to_be_between",
                "expect_column_median_to_be_between",
                "expect_column_quantile_values_to_be_between",
                "expect_column_values_to_be_in_set",
                "expect_column_proportion_of_unique_values_to_be_between",
                "expect_column_values_to_be_null",
            ]

        profiler = UserConfigurableProfiler(
            profile_dataset=validator,
            excluded_expectations=self.__excluded_expectations,
            ignored_columns=final_ignored_columns,
            not_null_only=False,
            primary_or_compound_key=False,
            semantic_types_dict=None,
            table_expectations_only=False,
            value_set_threshold="unique",
        )

        # this will build out a template, that you should tweak to your needs, in your great_expectations/expectations dir as {self.__suite}.json
        profiler.build_suite()
        validator.save_expectation_suite(discard_failed_expectations=False)
        self.__logger.info(
            f"Profiler found {len(lists_list)} List dtype values in the batch\n {lists_list}"
        )
        self.__logger.info(
            f"Profiler found {len(dicts_list)} Dict dtype values in the batch\n {dicts_list}"
        )
        self.__logger.info(
            f"Profiler will ignore {len(final_ignored_columns)} columns during the profiler\n {final_ignored_columns}"
        )

    @log_it
    def profiler_combine_to_json(
        self, days_back: int, profile_bucket_name: str, profile_prefix: str
    ):
        """
        this function will:
        1) check exists/create a ge_to_profile dir
        2) identify as many most recent files per day for how many days_back days is set in s3
        3) extend those json files into one larger json and save it within ge_to_profile
        In order for the profiler to analyze and assume thresholds for the template it builds, you need to feed it a json file
        containing a sampling of multiple records with the same columns.  If you care about averages, ranges, and NULLs then it is
        recommended to set days_back kwarg to multiple days back to increase accuracy.  If you know what your thresholds will be and do
        not want the recommended outcomes of the profiler, simply run a single day and manually add those values yourself.
        """
        if not os.path.isdir("ge_to_profile"):
            prompt_resp = (
                input(
                    f"\nge_to_profile does not exist in {os.getcwd()}\nwould you like to create this directory?\n[ y, n ]"
                )
                .lower()
                .strip()
            )

            if prompt_resp == "y":
                path = os.path.join(os.getcwd(), "ge_to_profile")
                # Create the directory
                os.mkdir(path)
        assert os.path.isdir(
            "ge_to_profile"
        ), "ge_to_profile directory must exist before running the profiler"

        month_list = wr.s3.list_directories(f"s3://{profile_bucket_name}/{profile_prefix}")
        days = []
        n = days_back
        # reverse the list of month= dirs and begin finding most recent days back
        for month in sorted(month_list, reverse=True):
            day_list = sorted(wr.s3.list_directories(month), reverse=True)
            # only extend by n as n changes on the available days
            days.extend(day_list[:n])
            # if there are less days than what is needed, reduce n and find more, else break the loop
            n -= len(day_list) if len(day_list) < n else n
            if len(days) == days_back:
                break
        to_profile = []
        # iterate all the identified days and make final json to profile out of most recent file per day
        for path in days:
            day = sorted(wr.s3.list_objects(path, ignore_empty=True), reverse=True)[0]
            df = wr.s3.read_json(day)
            to_profile.extend(df.to_dict(orient="records"))

        self.__logger.info(
            f"writing file ge_to_profile/{self.__data_name}_to_profile.json with {len(to_profile)} records"
        )
        with open(f"ge_to_profile/{self.__data_name}_to_profile.json", "w") as f:
            f.write(json.dumps(to_profile))
        self.__logger.info(f"ge_to_profile/{self.__data_name}_to_profile.json written")
