import sys
import time
import calendar
from datetime import datetime
import json

from elasticsearch.helpers import parallel_bulk

from utils.logger import Logger
from utils.pretty import Pretty
from utils.dates import Dates

from connectors.elasticsearch import elasticsearch_client
from connectors.snowflake import snowflake_cursor

from conf import config

def fetch_data_from_source(offset_gen):
    offset = offset_gen['offset']
    start_date_str = offset_gen['start_date_str']
    end_date_str = offset_gen['end_date_str']

    limit = config.INGEST_CONFIG['chunk_limit']
    try:
        snowflake_cursor.execute(f"""
            WITH top_sound AS (
                SELECT stp.*, pf.followers
                FROM STREAMS_BY_TRACK_PLAYLIST_COUNTRY_FEED_DISTRIBUTOR_DAILY AS stp
                LEFT JOIN MAPPINGS_DIM_PLAYLIST_TO_FOLLOWERS AS pf USING(stp.playlist_id)
                WHERE stp.DOWNLOAD_ACTIVITY_DATE between '{start_date_str}' and '{end_date_str}'
                LIMIT {limit}
                OFFSET {offset}
            )
            SELECT
                OBJECT_CONSTRUCT(
                    'isrc', ts.isrc,
                    'feed_id',ts.feed_id,
                    'store_id',ts.store_id,
                    'label_id',ts.label_id,
                    'track_id',ts.track_id,
                    'followers',ts.followers,
                    'product_id',ts.product_id,
                    'country_code',ts.country_code,
                    'playlist_id',ts.playlist_id,
                    'download_activity_date',ts.download_activity_date,
                    'store_playlist_id',ts.store_playlist_id,
                    'playlist_name',ts.playlist_name,
                    'playlist_url',ts.playlist_url,
                    'playlist_image',ts.playlist_image,
                    'streams',ts.streams
                )
            FROM top_sound as ts
        """)

        all_rows = snowflake_cursor.fetchall()

        return all_rows
    except Exception as error:
        Logger(f'{offset}: fetch_data_from_source_exception').log_exception(error)

def ingest_to_es(offset_gen):
    offset = offset_gen['offset']

    THREAD_COUNT = config.ES_BULK_CONFIG['thread_count']
    QUEUE_SIZE = config.ES_BULK_CONFIG['queue_size']
    MAX_CHUNK_BYTES = config.ES_BULK_CONFIG['max_chunk_bytes']
    CHUNK_SIZE = config.ES_BULK_CONFIG['chunk_size']

    Logger(f'{offset}: START: ingest_to_es').log_info()

    try:
        start = time.time()
        Logger(f'{offset}: Before Query Current Timestamp').log_info()

        all_rows = fetch_data_from_source(offset_gen)

        query_seconds = time.time() - start

        Logger(f'{offset}: It took for query {str(query_seconds)} seconds for {len(all_rows)} rows.').log_info()
        Logger(f'{offset}: After Query & Before Ingestion').log_info()

        def gen_data():
            for row in all_rows:
                es_source_json = json.loads(row[0])
                parsed_date = datetime.strptime(es_source_json['download_activity_date'], '%Y-%m-%d')

                index_date = f'{parsed_date.year}-{parsed_date.month}'
                index_name = f'{config.ES_INDEX}-{index_date}'

                yield {
                    '_index': index_name,
                    '_type': '_doc',
                    '_source': es_source_json,
                }

        start = time.time()
        for success, info in parallel_bulk(elasticsearch_client, gen_data(), thread_count=THREAD_COUNT, \
            queue_size=QUEUE_SIZE, max_chunk_bytes=MAX_CHUNK_BYTES, chunk_size=CHUNK_SIZE):
            if not success:
                error = Pretty(info).print()
                Logger(f'{offset}: A document failed: {error}').log_info()

        seconds = time.time() - start

        full_push_time = query_seconds + seconds
        Logger(f'{offset}: it took for parallel_bulk {str(seconds)} seconds.').log_info()
        Logger(f'{offset}: total time {full_push_time}').log_info()

        Logger(f'{offset}: After Ingestion').log_info()

        Logger(f'{offset}: END: ingest_to_es').log_info()

    except Exception as error:
        Logger(f'{offset}: ingest_to_es_exception').log_exception(error)

def clone_index(index_name, remove=False):
    cloned_index = f'cloned-{index_name}'

    try:
        res = elasticsearch_client.indices.delete(index=cloned_index, ignore_unavailable=True)
        Logger(f'\t{res} Deletion of cloned index : {cloned_index}').log_info()
    except Exception as error:
        Logger(f'delete_cloned_indexes_exception {cloned_index}').log_exception(error)

    if remove is True:
        return

    try:
        res = elasticsearch_client.indices.put_settings(
            index=index_name,
            body={
                'index': {
                    'blocks.write': True
                }
            })
        Logger(f'\t{res} blocks.write index : {index_name}').log_info()
    except Exception as error:
        Logger(f'blocks.write_indexes_exception {index_name}').log_exception(error)

    try:
        res = elasticsearch_client.indices.clone(index=index_name, target=cloned_index)
        Logger(f'\t{res} Cloned index : {cloned_index}').log_info()

        res = elasticsearch_client.indices.put_settings(
            index=cloned_index,
            body={
                'index': {
                    'blocks.write': None
                }
            })
        Logger(f'\t{res} Reset blocks.write cloned index : {cloned_index}').log_info()
    except Exception as error:
        Logger(f'clone_indexes_exception {index_name}').log_exception(error)

    try:
        res = elasticsearch_client.indices.put_alias(index=cloned_index, \
            name=config.ES_INDEX_ALIAS)
        Logger(f'\t{res} Set alias for cloned index : {cloned_index}').log_info()
    except Exception as error:
        Logger(f'delete_indexes_exception {index_name}').log_exception(error)

def recreate_index(index_name):
    Logger(f'START: recreate_index: {index_name}').log_info()
    try:
        res = elasticsearch_client.indices.delete(index=index_name, ignore_unavailable=True)
        Logger(f'\t{res} Deletion of index : {index_name}').log_info()
    except Exception as error:
        Logger(f'delete_indexes_exception {index_name}').log_exception(error)

    try:
        res = elasticsearch_client.indices.create(
            index=index_name,
            body={
                'settings': config.ES_IDX_INGEST_SETTINGS
            })
        Logger(f'\t{res} Creation of index : {index_name}').log_info()
    except Exception as error:
        Logger(f'create_indexes_exception {index_name}').log_exception(error)

    Logger(f'END: recreate_index: {index_name}').log_info()

def prepare_index_to_be_ingested(index_name):
    Logger('START: prepare_index_to_be_ingested').log_info()

    clone_index(index_name)

    recreate_index(index_name)

    Logger('END: prepare_index_to_be_ingested').log_info()

def fetch_data_count_from_snowflake(start_date_str, end_date_str):
    snowflake_cursor.execute(f"""
        SELECT COUNT(*)
        FROM STREAMS_BY_TRACK_PLAYLIST_COUNTRY_FEED_DISTRIBUTOR_DAILY AS stp
        WHERE stp.DOWNLOAD_ACTIVITY_DATE BETWEEN '{start_date_str}' AND '{end_date_str}'
    """)

    snowflake_row_count = snowflake_cursor.fetchone()[0]
    Logger(f'Total rows in snowflake: {snowflake_row_count}').log_info()

    return snowflake_row_count

def reset_index_settings():
    index_name = f'{config.ES_INDEX}-*'
    try:
        res = elasticsearch_client.indices.put_settings(
            index=index_name,
            body={
                'index': config.ES_IDX_SEARCH_SETTINGS,
            })

        Logger(f'{res} Reset index for query : {index_name}').log_info()

    except Exception as error:
        Logger('reset_indexes_exception').log_exception(error)

    try:
        res = elasticsearch_client.indices.put_alias(
            index=index_name,
            name=config.ES_INDEX_ALIAS
            )

        Logger(f'{res} Set alias for indexes : {index_name}').log_info()

    except Exception as error:
        Logger('set_alias_exception').log_exception(error)