import gzip
import os
from io import BytesIO

from cameo_claw.cameo_claw_reuse import bytesio_to_gzip, requests_get_write, \
    bytesio_filename, write_gzip, pl_read_csv_distinct, pd_read_csv
from cameo_claw.functional import it_mp_f
import polars as pl
import pandas as pd


def a_download(url, target_directory):
    def write(bytes1):
        filename = os.path.basename(url)
        path = f'{target_directory}{filename}'
        with open(path, 'wb') as f:
            f.write(bytes1)
        return url

    return requests_get_write(target_directory, url, write)


def it_download(lst_url, target_directory):
    return it_mp_f(a_download, [tuple([url, target_directory]) for url in lst_url])


def a_distinct(url, target_directory, lst_distinct_column):
    def write(bytes1):
        df = pl_read_csv_distinct(bytes1, lst_distinct_column)
        return bytesio_to_gzip(df, url, target_directory)

    return requests_get_write(target_directory, url, write)


def it_distinct(lst_url, target_directory, lst_distinct_column):
    return it_mp_f(a_distinct, [tuple([url, target_directory, lst_distinct_column]) for url in lst_url])


def a_group(url, target_directory, lst_distinct_column, lst_group_by_column):
    def write(bytes1):
        df = pl_read_csv_distinct(bytes1, lst_distinct_column)
        g = df.groupby(lst_group_by_column)
        for df in g:
            filename_tail = '_group_' + '-'.join(list(
                map(lambda column:
                    str(df.row(0)[df.find_idx_by_name(column)]).replace('_', '-'),
                    lst_group_by_column)))
            bytesio, filename = bytesio_filename(df, url)
            path = f'{target_directory}{filename}{filename_tail}.csv.gz'
            write_gzip(bytesio, path)
        return url

    return requests_get_write(target_directory, url, write)


def it_group(lst_url, target_directory, lst_distinct_column, lst_group_by_column):
    return it_mp_f(a_group,
                   [tuple([url, target_directory, lst_distinct_column, lst_group_by_column]) for url in lst_url])


def a_filter(url, target_directory, lst_distinct_column, lst_column_match=[], sort_column=''):
    def write(bytes1):
        df = pl_read_csv_distinct(bytes1, lst_distinct_column)
        condition = ''
        for lst_col_val in lst_column_match:
            condition += f"""(pl.col('{lst_col_val[0]}')=='{lst_col_val[1]}') | """
        condition = condition[:-3]
        df = df.filter(eval(condition))
        if sort_column:
            df = df.sort(sort_column)
        bytesio_to_gzip(df, url, target_directory)
        return url

    return requests_get_write(target_directory, url, write)


def it_filter(lst_url, target_directory, lst_distinct_column, lst_column_match, sort_column):
    return it_mp_f(a_filter,
                   [tuple([url, target_directory, lst_distinct_column,
                           lst_column_match, sort_column]) for url in lst_url])


def a_chunk(url, target_directory, lst_distinct_column, lst_column_match=[], sort_column=''):
    def write(bytes1):
        # 2022-04-10 06:21
        # use polars hang at FastAPI+Python 3.7
        # use polars hang at FastAPI+replit Python 3.8.12
        # use polars OK at FastAPI+Mac Python 3.8.9
        df = pd_read_csv(bytes1)
        df = df.drop_duplicates(subset=lst_distinct_column)

        # pandas column equals example
        # rslt_df = dataframe[(dataframe['Age'] == 22) &
        #                     dataframe['Stream'].isin(options)]
        condition = ''
        for lst_col_val in lst_column_match:
            condition += f"""(df['{lst_col_val[0]}']=='{lst_col_val[1]}') | """
        condition = condition[:-3]
        df = df[eval(condition)]

        if sort_column:
            df = df.sort_values(by=[sort_column])
        bytesio = BytesIO()
        df.to_csv(bytesio, header=False, index=False, compression='gzip')
        return bytesio.getvalue()

    return requests_get_write(target_directory, url, write)


def it_chunk_pd(lst_url, target_directory, lst_distinct_column, lst_column_match, sort_column):
    return it_mp_f(a_chunk,
                   [tuple([url, target_directory, lst_distinct_column,
                           lst_column_match, sort_column]) for url in lst_url])


def a_chunk_pl(url, target_directory, lst_distinct_column, lst_column_match=[], sort_column=''):
    def write(bytes1):
        df = pl.read_csv(bytes1)
        df = df.distinct(subset=lst_distinct_column)
        condition = ''
        for lst_col_val in lst_column_match:
            condition += f"""(pl.col('{lst_col_val[0]}')=='{lst_col_val[1]}') | """
        condition = condition[:-3]
        df = df.filter(eval(condition))
        if sort_column:
            df = df.sort(sort_column)
        bytesio_csv = BytesIO()
        df.to_csv(bytesio_csv, has_header=False)
        bytesio_gzip = BytesIO()
        with gzip.open(bytesio_gzip, 'wb') as f:
            f.write(bytesio_csv.getvalue())
        return bytesio_gzip.getvalue()

    return requests_get_write(target_directory, url, write)


def a_chunk_pl_lst(url, target_directory, lst_distinct_column, lst_column_match, sort_column):
    tup_distinct_column = tuple(lst_distinct_column)
    tup_lst_column_match = tuple([tuple(l) for l in lst_column_match])
    return a_chunk_pl_tup(url, target_directory, tup_distinct_column, tup_lst_column_match, sort_column)


from functools import lru_cache


@lru_cache(maxsize=5000)
def a_chunk_pl_tup(url, target_directory, tup_distinct_column, tup_column_match, sort_column):
    # print('no cache:', url, target_directory, tup_distinct_column, tup_column_match, sort_column)

    def write(bytes1):
        df = pl.read_csv(bytes1)
        df = df.distinct(subset=list(tup_distinct_column))
        condition = ''
        for lst_col_val in tup_column_match:
            condition += f"""(pl.col('{lst_col_val[0]}')=='{lst_col_val[1]}') | """
        condition = condition[:-3]
        df = df.filter(eval(condition))
        if sort_column:
            df = df.sort(sort_column)
        bytesio_csv = BytesIO()
        df.to_csv(bytesio_csv, has_header=False)
        bytesio_gzip = BytesIO()
        with gzip.open(bytesio_gzip, 'wb') as f:
            f.write(bytesio_csv.getvalue())
        return bytesio_gzip.getvalue()

    return requests_get_write(target_directory, url, write)


def it_chunk_pl(lst_url, target_directory, lst_distinct_column, lst_column_match, sort_column):
    return it_mp_f(a_chunk_pl_lst,
                   [tuple([url, target_directory, lst_distinct_column,
                           lst_column_match, sort_column]) for url in lst_url])


def a_filter_group(url, target_directory, lst_distinct_column, lst_column_match=[], sort_column='',
                   lst_group_by_column=[]):
    def write(bytes1):
        df = pl_read_csv_distinct(bytes1, lst_distinct_column)
        condition = ''
        for lst_col_val in lst_column_match:
            condition += f"""(pl.col('{lst_col_val[0]}')=='{lst_col_val[1]}') | """
        condition = condition[:-3]
        df = df.filter(eval(condition))
        if sort_column:
            df = df.sort(sort_column)
        g = df.groupby(lst_group_by_column)
        for df in g:
            filename_tail = '_group_' + '-'.join(list(
                map(lambda column:
                    str(df.row(0)[df.find_idx_by_name(column)]).replace('_', '-'),
                    lst_group_by_column)))
            bytesio, filename = bytesio_filename(df, url)
            path = f'{target_directory}{filename}{filename_tail}.csv.gz'
            write_gzip(bytesio, path)
        return url

    return requests_get_write(target_directory, url, write)


def it_filter_group(lst_url, target_directory, lst_distinct_column, lst_column_match, sort_column, lst_group_by_column):
    return it_mp_f(a_filter_group,
                   [tuple([url, target_directory, lst_distinct_column,
                           lst_column_match, sort_column, lst_group_by_column]) for url in lst_url])
