import os
from typing import List, Any
import logging
import boto3
import awswrangler as wr
import json

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: str,
        prefix: str,
        excluded_expectations: list = [],
        logger: logging.Logger = None,
        aws_session: Any = None,
    ):
        self.__data_name = data_name
        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
        )

        #TODO: keep commented for when it is time to address simplifying ignoring columns
        # column_names = [column_name for column_name in validator.columns()]
        # validator.head(n_rows=5, fetch_all=False)
        # ignored_columns = column_names
        
        #TODO: manually edited ignore columns for demo/sample set, needs to be more dynamic or more control
        # ignored_columns.remove('engaged_seconds')
        # ignored_columns.remove('engaged_user_count')
        # ignored_columns.remove('escalation_count')
        # ignored_columns.remove('major')
        # ignored_columns.remove('seconds_to_engage')
        # ignored_columns.remove('seconds_to_first_ack')
        # ignored_columns.remove('seconds_to_mobilize')
        # ignored_columns.remove('seconds_to_resolve')
        # ignored_columns.remove('user_defined_effort_seconds')

        # 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"
            ]

        profiler = UserConfigurableProfiler(
            profile_dataset=validator,
            excluded_expectations=self.__excluded_expectations,
            ignored_columns= None, #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)
    
    @log_it
    def profiler_combine_to_json(self, days_back: int):
        '''
        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(
            self.__s3_arn
        )
        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')
