import logging
import os
import time
from sys import platform
from typing import List

import pandas as pd

from service.async_task_manager import ThreadWorkerQueue, async_task
from service.tasks.airflow.enrichment import AirflowEnrichmentRunner
from service.tasks.api_service_handler.klaviyo.service_klaviyo import KlaviyoServiceCore
from service.tasks.api_service_handler.klaviyo.service_klaviyo_db_handler import (
    DbKlaviyoServiceHandler,
)
from service.tasks.api_service_handler.utils import chunks, field_mapper
from service.tasks.audience.audience_klaviyo import KlaviyoAudience
from service.tasks.base_attributes_generation import (
    generate_labels_and_attributes,
    process_file,
    update_collection_labels,
)
from service.tasks.data_check.field_map_checker import FieldMapChecker
from service.tasks.emailing.email_alliance import send_alliance_email
from service.tasks.emailing.email_external_integration import (
    send_export_confirmation_email,
)
from service.tasks.emailing.email_subscription import send_subscription_email
from service.tasks.enrichment import EnrichKlaviyoEventsData, EnrichKlaviyoStaticData
from service.tasks.filtering import save_segment
from service.tasks.guess_system_fields import guess_fields
from service.tasks.management.alliance import (
    generate_alliance_segments,
    get_alliance,
    merge_collections_into_alliance,
)
from service.tasks.precondition.checks.deletion import (
    AllianceCollectionDeleteCheck,
    AllianceDeleteCheck,
    AllianceMemberRemoveCheck,
    WorkspaceCollectionDeleteCheck,
    WorkspaceDeleteCheck,
    WorkspaceSegmentByIdDeleteCheck,
)
from service.tasks.precondition.engine import Precondition
from service.utils import appsync_communication as appsync
from service.utils.appsync_communication.message_types import ServiceInfoError
from service.utils.aws_connectors import run_query
from service.utils.data_model_utils import (
    change_collection_source,
    change_collection_status,
    delete_used_guesses,
    generate_collection,
    get_related_collections,
    get_schema_collection_id_for_url_path,
    get_source_from_collection,
    read_saved_field_map,
    save_field_map,
    update_collection,
)
from service.utils.s3_file_operations import s3_move

logger = logging.getLogger(__name__)

# Initialize thread workers
THREAD_COUNT = int(os.getenv("ASYNC_WORKERS", 8))
IO_THREAD_COUNT = int(os.getenv("IO_WORKERS", 100))
ThreadWorkerQueue.instantiate(
    thread_count=THREAD_COUNT, io_worker_limit=IO_THREAD_COUNT
)
"""Can't run parallel_apply on windows machines"""
if platform == "linux":
    from pandarallel import pandarallel

    pandarallel.initialize()


@async_task(async_source="aws:s3.ObjectCreated:Put")
def handle_s3_event(data, bucket, key, event_name, task_handle=None, **kw):
    """Acknowledge file upload"""
    try:
        bucket_name = bucket
        file_key = key
        s3_event_name = event_name
    except Exception as e:
        if task_handle:
            task_handle.send_immediate_error("validation fail: invalid parameters")
        logger.exception(e)
        raise RuntimeError("Did not find expected fields in s3 event record")

    # validate s3 event:
    if (
        s3_event_name not in ["ObjectCreated:Put"]
        or not bucket_name
        or not file_key
        or not s3_event_name
    ):
        if task_handle:
            task_handle.send_immediate_error(
                f"validation fail: Not enough data to handle s3 event: {data}"
            )
        error = RuntimeError(f"Not enough data to handle s3 event: {data}")
        logger.exception(error)
        raise error
    # initial checks passed, return immediate response.
    if task_handle:
        task_handle.send_immediate_response({"status": "accepted"})
    # get collection id for file key
    try:
        (
            workspace_schema,
            collection_id,
            user_id,
        ) = get_schema_collection_id_for_url_path(file_key)[0]
    except Exception as e:
        logger.exception(e)
        raise RuntimeError(
            "Didn't find upload in progress from commons.upload_url table."
        )

    # rename the file into something more nice
    old_file_key = file_key
    nice_file_key = f"{workspace_schema}/uploads/{file_key.replace('/','_')}"
    try:
        s3_move(bucket_name, file_key, bucket_name, nice_file_key)
        change_collection_source(
            workspace_schema, collection_id, f"{bucket_name}/{nice_file_key}"
        )
        file_key = nice_file_key
    except Exception as e:
        logger.exception(e)
    if file_key == nice_file_key:
        logger.info(f"Renamed [{bucket_name}]: {old_file_key} -> {file_key}")
    else:
        raise RuntimeError(
            f"S3 rename file failed! [{bucket_name}]: {old_file_key} -> {file_key}"
        )

    # set uploaded status
    collection_status = "uploaded"
    change_collection_status(
        workspace_schema, user_id, collection_id, collection_status
    )

    # proceed with labeling
    try:
        response = guess_fields(
            collection_id, workspace_schema, bucket=bucket_name, file_name=file_key
        )
        if response["status"] == "success":
            collection_status = "labeling edit"
        else:
            collection_status = "initial checking fail"
    except Exception as e:
        logger.exception(e)
        collection_status = "initial checking fail"

    change_collection_status(
        workspace_schema, user_id, collection_id, collection_status
    )


@async_task(async_source="appsync.mapCollection")
def handle_process_file(
    data,
    workspace_schema,
    collectionId,
    fieldMap,
    user_id,
    collectionName=None,
    task_handle=None,
    management_schema=None,
    **kw,
):
    """Stores labels, and processes csv file to unpack data into data model
    based on tasks_controller.process_scv_file, but takes collectionId instead of fileId
    """
    checker = FieldMapChecker(data=fieldMap)
    check_results = checker.verify()
    if not check_results["check_passed"]:
        if task_handle:
            task_handle.send_immediate_error(
                f"validation fail: {check_results['message']}"
            )
        raise RuntimeError(
            f'Did not find expected fields in fieldMap - {check_results["message"]}'
        )

    start_time = time.time()

    """ Save the field map on top of guesses """
    save_field_map(workspace_schema, collectionId, fieldMap)

    """ Set status and rename collection immediately """
    try:
        collection = update_collection(
            workspace_schema,
            user_id,
            collectionId,
            name=collectionName,
            status="labeling done",
        )
        if task_handle:
            task_handle.send_immediate_response(collection)
    except Exception as e:
        if task_handle:
            task_handle.send_immediate_error(e)
        logger.exception(e)
        raise e

    source = get_source_from_collection(workspace_schema, collectionId)

    if source.startswith("klaviyo"):
        """Working with previously imported data from 3rd party integration. Klaviyo service."""
        """Getting assigned system fields with system_field_ids"""

        field_map_df = read_saved_field_map(workspace_schema, collectionId)

        collection_ids = get_related_collections(
            schema=workspace_schema, collection_id=collectionId
        )

        """Separating attributes that we need to transfer from system fields to none system fields"""
        cond = field_map_df["system_field_name"] == ""
        stn_field_map = field_map_df[cond]

        if len(stn_field_map) > 0:
            """Making unique field names to use as file_field_name will be a valid system field name here"""
            stn_field_map["system_field_name"] = stn_field_map["file_field_name"].apply(
                lambda x: str(x) + "_"
            )
            existing_attributes = run_query(
                f"SELECT id, name AS system_field_name FROM {workspace_schema}.attribute;",
                return_type="df",
            )

            """Checking if we already done the same in the past and field exists with it's id"""
            stn_field_map = stn_field_map.merge(
                existing_attributes, how="left", on="system_field_name"
            )
            cond2 = stn_field_map["id"].isna()

            """Using existing attribute id for existing non system fields"""
            stn_field_map_a = stn_field_map[~cond2]
            stn_field_map_a["system_field_id"] = stn_field_map_a["id"]

            """Generating new attirbute id for non esisint non system fields"""
            stn_field_map_b = stn_field_map[cond2]
            next_id = int(existing_attributes.id.max()) + 1
            stn_field_map_b["system_field_id"] = range(
                next_id, next_id + len(stn_field_map_b)
            )

            stn_field_map = pd.concat(
                [stn_field_map_a, stn_field_map_b], ignore_index=True
            )
            stn_field_map.drop(columns="id", inplace=True)

            stn_field_map = stn_field_map.to_dict("records")
            update_collection_labels(
                schema_name=workspace_schema,
                collection_ids=collection_ids,
                field_map=stn_field_map,
                to_system_fields=False,
            )

        """Separating attributes that we need to transfer from none system fields to system fields"""
        nts_field_map = field_map_df[~cond]
        nts_field_map = nts_field_map.to_dict("records")

        update_collection_labels(
            schema_name=workspace_schema,
            collection_ids=collection_ids,
            field_map=nts_field_map,
        )

        """Remove guesses from guessed_file_uplead_fields table"""
        collection = update_collection(
            workspace_schema,
            user_id,
            collectionId,
            name=collectionName,
            status="enriching",
        )
        delete_used_guesses(workspace_schema, collectionId)
        collection_id = collectionId
    else:
        source = source.split("/")
        bucket_name, file_name = source[0], "/".join(source[1:])

        # Renaming already done, now the collection_rename inside process_file doesn't do anything, if new name is None.
        collection_id = process_file(
            workspace_schema,
            management_schema,
            file_name,
            bucket_name,
            collectionName,
            fieldMap,
            user_id,
            collectionId,
            final_status="processing",
        )

    airflow_runner = AirflowEnrichmentRunner(
        user_id=user_id, workspace_schema=workspace_schema, collection_id=collection_id
    )
    try:
        airflow_runner.execute_collection_enrichment()
        change_collection_status(workspace_schema, user_id, collection_id, "enriching")
    except Exception as e:
        logger.error(e, exc_info=e)
        change_collection_status(workspace_schema, user_id, collection_id, "failed")
        raise RuntimeError(f"Airflow run has failed to start: {e}")

    elapsed_time = time.time() - start_time
    logger.info(
        f"{workspace_schema} - successfully sent source collection {collection_id} enrichment task to Airflow in {elapsed_time} seconds"
    )


@async_task(async_source="appsync.deleteCollection")
def handle_delete_collection(
    data,
    workspace_schema,
    management_schema,
    user_id,
    collectionId=None,
    collectionIds=None,
    acceptedConsequences=None,
    task_handle=None,
    **kw,
):
    """
    Removes all the data related to the collection_id
    """
    # logger.info(f"input payload: {repr(data)}")
    precondition = Precondition(
        WorkspaceCollectionDeleteCheck(
            workspace_schema,
            collectionIds or collectionId,
            management_schema=management_schema,
            user_id=user_id,
        ),
        WorkspaceSegmentByIdDeleteCheck(
            workspace_schema,
            collectionIds or collectionId,
            management_schema=management_schema,
            user_id=user_id,
        ),
        accepted_consequences=acceptedConsequences,
    )
    if precondition.is_met:
        collection = change_collection_status(
            workspace_schema, user_id, collectionId, "deleting"
        )
        task_handle.send_immediate_response(collection)
        precondition.execute()
    else:
        task_handle.send_immediate_error(precondition.get_exception())


@async_task(async_source="appsync.sendAllianceMemberEmail")
def send_invite_to_alliance_member(
    data,
    template_type,
    inviter_email,
    recipient,
    alliance_name,
    message=None,
    temp_passwd=None,
    task_handle=None,
    **kw,
):
    """
    Sends emails to new or existing alliance members. Since it contains PII (sender/recipient email),
    we don't log the payload like we do in other controller functions.

    templates types:
        - invite_new_user
        - invite_existing_user
        - remove_member
    """
    if template_type == "invite_new_user" and not temp_passwd:
        raise Exception(
            f'Did not get required "temp_passwd" for {template_type} alliance template'
        )

    if template_type in ("invite_new_user", "invite_existing_user") and not message:
        message = ""

    task_handle.send_immediate_response("sending")
    send_alliance_email(
        template_type, inviter_email, recipient, alliance_name, message, temp_passwd
    )


@async_task(async_source="appsync.sendSubscriptionEmail")
def send_subscription_email_async(
    data,
    template_type,
    recipient,
    package_name,
    billing_cycle=None,
    active_date=None,
    next_payment_date=None,
    task_handle=None,
    **kw,
):
    """
    Sends emails to new or existing alliance members. Since it contains PII (sender/recipient email),
    we don't log the payload like we do in other controller functions.

    templates types:
        - invite_new_user
        - invite_existing_user
        - remove_member
    """
    task_handle.send_immediate_response("sending")
    send_subscription_email(
        template_type,
        recipient,
        package_name,
        billing_cycle,
        active_date,
        next_payment_date,
    )


@async_task(async_source="appsync.saveSegment")
def get_save_segment(
    data,
    setId,
    user_id,
    workspace_schema=None,
    alliance_schema=None,
    name=None,
    task_handle=None,
    **kw,
):
    """save_segment
    all possible, if no set_id
    only applicable to a collection, if existing set_id points to source collection
    saved pre-filled filters, if existing set_id points to set/segment"""
    schema = alliance_schema or workspace_schema
    collection = save_segment(
        schema=schema, user_id=user_id, set_id=setId, new_name=name
    )
    task_handle.send_immediate_response(collection)

    airflow_runner = AirflowEnrichmentRunner(
        user_id=user_id,
        workspace_schema=workspace_schema,
        alliance_schema=alliance_schema,
        collection_id=setId,
    )
    try:
        airflow_runner.execute_segment_enrichment()
        change_collection_status(schema, user_id, setId, "enriching")
    except Exception as e:
        logger.error(e, exc_info=e)
        change_collection_status(schema, user_id, setId, "failed")
        raise RuntimeError(f"Airflow run has failed to start: {e}")

    logger.info(
        f"{schema} - successfully sent collection {setId} segment generation task to Airflow!"
    )


@async_task(async_source="appsync.mergeProfilesIntoAlliance")
def merge_profiles_to_alliance(
    data,
    workspace_schema,
    user_id,
    alliance_schema,
    input,
    management_schema,
    task_handle=None,
    **kw,
):
    # validate inputs
    merge_collection_ids = []
    workspace_set = set()
    for coll in input:
        cs = coll.split("-")
        if len(cs) == 2:
            merge_collection_ids.append(str(int(cs[-1])))
            workspace_set.add(cs[0])
        else:
            raise RuntimeError("Invalid id")
    if len(workspace_set) == 1:
        for ws in run_query(
            f"SELECT id FROM {management_schema}.workspace WHERE id = %(workspace)s",
            {"workspace": list(workspace_set)[0]},
        ):
            workspace_schema = ws[0]
            break
        else:
            raise RuntimeError("Invalid id")
    else:
        raise RuntimeError("Cant merge collections from different workspaces at once")

    related_collection_ids = get_related_collections(
        workspace_schema, merge_collection_ids, only_types=["source", "enrichment"]
    )
    if related_collection_ids:
        # now do the actual thing, giving task handle in,
        # so the first update would already have all the collections, just empty.
        merge_collections_into_alliance(
            workspace_schema,
            management_schema,
            related_collection_ids,
            alliance_schema,
            user_id,
            task_handle,
        )
        appsync.send_alliance(
            alliance_schema, {"id": alliance_schema, "status": "updating"}
        )
        generate_alliance_segments(alliance_schema, user_id)
        appsync.send_alliance(alliance_schema, get_alliance(alliance_schema))
    else:
        raise RuntimeError("Invalid id")


@async_task(async_source="appsync.removeAllianceMember")
def remove_alliance_member(
    data,
    alliance_schema,
    user_id,
    management_schema,
    member_schema,
    member_user_id,
    alliance_management_schema,
    acceptedConsequences,
    task_handle=None,
    **kw,
):
    precondition = Precondition(
        AllianceMemberRemoveCheck(
            alliance_schema,
            member_user_id,
            management_schema=alliance_management_schema,
            user_id=user_id,
        ),
        accepted_consequences=acceptedConsequences,
    )

    if precondition.is_met:
        task_handle.send_immediate_response(True)
        precondition.execute()
    else:
        task_handle.send_immediate_error(precondition.get_exception())


@async_task(async_source="appsync.deleteAllianceCollections")
def delete_alliance_collections(
    data,
    alliance_schema,
    management_schema,
    user_id,
    collection_ids,
    task_handle=None,
    acceptedConsequences=None,
    **kw,
):
    precondition = Precondition(
        AllianceCollectionDeleteCheck(
            alliance_schema,
            collection_ids,
            management_schema=management_schema,
            user_id=user_id,
        ),
        accepted_consequences=acceptedConsequences,
    )
    if precondition.is_met:
        task_handle.send_immediate_response(True)
        precondition.execute()
    else:
        task_handle.send_immediate_error(precondition.get_exception())


@async_task(async_source="appsync.deleteWorkspace")
def handle_delete_workspace(
    data,
    user_id,
    workspace_schema,
    management_schema,
    acceptedConsequences=None,
    task_handle=None,
    **kw,
):
    precondition = Precondition(
        WorkspaceDeleteCheck(
            workspace_schema, management_schema=management_schema, user_id=user_id
        ),
        accepted_consequences=acceptedConsequences,
    )

    if precondition.is_met:
        task_handle.send_immediate_response(True)
        precondition.execute()
    else:
        task_handle.send_immediate_error(precondition.get_exception())


@async_task(async_source="appsync.deleteAlliance")
def handle_delete_alliance(
    data,
    user_id,
    alliance_schema,
    alliance_management_schema,
    acceptedConsequences=None,
    task_handle=None,
    **kw,
):
    precondition = Precondition(
        AllianceDeleteCheck(
            alliance_schema,
            management_schema=alliance_management_schema,
            user_id=user_id,
        ),
        accepted_consequences=acceptedConsequences,
    )

    if precondition.is_met:
        task_handle.send_immediate_response(True)
        precondition.execute()
    else:
        task_handle.send_immediate_error(precondition.get_exception())


@async_task(async_source="appsync.klaviyoImport")
@field_mapper
def get_klaviyo_profile_base_data_from_list(
    data,
    management_schema,
    workspace_schema,
    user_id,
    collectionName,
    listIds: List[str],
    task_handle=None,
    **kw,
):
    start_time = time.time()
    kdb = DbKlaviyoServiceHandler()
    confirmed_df = kdb.confirm_klaviyo_list_ids(
        schema=management_schema, list_ids=listIds
    )
    if (len(confirmed_df) == 0) or (len(confirmed_df) != len(listIds)):
        e = ServiceInfoError("Wrong listIds provided. Please provide correct listIds.")
        task_handle.send_immediate_error(e)
        logger.exception(e)
        raise e
    else:
        pass

    klaviyo = KlaviyoServiceCore(schema=management_schema)
    data_list = []
    for lst in listIds:
        response = klaviyo.get_list_or_segment_persons(list_id=lst)
        for r in response:
            df = pd.DataFrame(r["data"]["records"])
            data_list.append(df)

    """Combining results from all list_ids and removing duplicates. (same person can be in multiple list_ids)"""
    source_df = pd.concat(data_list, ignore_index=True)
    source_df.drop_duplicates(inplace=True, ignore_index=True)

    """Formatting the data and ensuring same structure every time:"""
    if len({"id", "email"} - set(source_df.columns)) == 0:
        source_df = source_df[["email", "id"]]
    else:
        raise RuntimeError(
            "Unexpected response from Klaviyo group/{list_id}/members/all endpoint"
        )

    collection_source = f'klaviyo: {"|".join(listIds)}'

    """Generating collection_id"""
    collection_id = generate_collection(
        collection_name=f"{collectionName}",
        collection_source=f"{collection_source}",
        schema_name=workspace_schema,
        user_id=user_id,
        status="downloading",
    )

    try:
        collection = update_collection(
            workspace_schema,
            user_id,
            collection_id,
            name=collectionName,
            status="downloading",
        )
        if task_handle:
            task_handle.send_immediate_response(collection)
    except Exception as e:
        if task_handle:
            task_handle.send_immediate_error(e)
        logger.exception(e)
        raise e

    """Expected format coming from Klaviyo API"""
    field_map = [
        {"value": "userEmail", "index": 0},
        {"value": "klaviyoPersonId", "index": 1},
    ]
    response = generate_labels_and_attributes(
        schema_name=workspace_schema,
        management_schema=management_schema,
        collection_name=f"{collectionName}",
        field_map=field_map,
        source_df=source_df,
        user_id=user_id,
        collection_id=collection_id,
        final_status="downloading",
    )

    try:
        for core_enr in [
            EnrichKlaviyoEventsData,
            EnrichKlaviyoStaticData,
        ]:  # StaticData should be executed after EventsData
            enr = core_enr(
                collection_id=collection_id,
                schema=workspace_schema,
                user_id=user_id,
                management_schema=management_schema,
            )
            enr.do_everything()
    except Exception as e:
        logger.exception(e)

    """Preparing attributes for labeling
       Ignoring some fields to clean the labeling screen for user and avoid issues with hidden system labels
    """
    fields_to_ignore = [
        "id",
        "statistic_id",
        "uuid",
        "klaviyoPersonId",
        "object",
        "event_properties.Method",
        "$title",
        "$organization",
        "$email",
        "email",
        "$id",
        "detail",
        "event_properties.$event_id",
        "timestamp",
    ]
    kdb.populate_guessed_file_upload_fields(
        schema=workspace_schema,
        collection_id=collection_id,
        ignore_system_fields=False,
        fields_to_ignore=fields_to_ignore,
    )
    """Setting up collection status to allow front end to start labeling process"""
    collection = update_collection(
        workspace_schema,
        user_id,
        collection_id,
        name=collectionName,
        status="labeling edit",
    )
    elapsed_time = time.time() - start_time
    logger.info(
        f"{workspace_schema} - successfully imported Klaviyo data for collection_id {collection_id} in {elapsed_time} seconds"
    )

    return "finished"


@async_task(async_source="appsync.klaviyoExport")
@field_mapper
def export_klaviyo_list(
    data,
    management_schema,
    email,
    workspace_schema,
    collectionId,
    alliance_schema=None,
    listName=None,
    task_handle=None,
    **kw,
):

    """List the available users to upload to Klaviyo"""
    audience = KlaviyoAudience(
        collection_id=collectionId,
        schema=workspace_schema,
        alliance_schema=alliance_schema,
    )
    users = audience.get_audience_data()
    users_dict = users.to_dict("records")

    if len(users_dict) < 1:
        msg = "You can only export to Klaviyo your owned first party fan profiles. This collection does not contain such fans."
        e = ServiceInfoError(msg)
        if task_handle:
            task_handle.send_immediate_error(e)
        logger.exception(e)
        raise e

    klaviyo = KlaviyoServiceCore(schema=management_schema)
    list_name = f"Fansifter - {listName}"
    list_creation_response = klaviyo.create_list(list_name=list_name)

    data = list_creation_response.get("data", {})

    if "list_id" in data:
        audience = KlaviyoAudience(
            collection_id=collectionId,
            schema=workspace_schema,
            alliance_schema=alliance_schema,
        )
        users = audience.get_audience_data()

        response_obj = {"status": "exporting", "profileCount": len(users)}
        if task_handle:
            task_handle.send_immediate_response(response_obj)

        users_dict = users.to_dict("records")

        response_list = []
        async_response_list = []
        for chunk in chunks(users_dict, 100):
            async_response_list.append(
                klaviyo.add_users_to_a_list(list_id=data["list_id"], users=chunk)
            )

        for task in async_response_list:
            response = task.wait_for_result()
            response_list.append(response)

        collection_data = audience.get_general_data()
        collection_data = collection_data.to_dict("records")

        send_export_confirmation_email(
            recipient=email,
            audience_name=collection_data[0].get("name", None),
            external_service_name="Klaviyo",
            number_of_fans=len(users_dict),
            external_list_name=list_name,
            external_list_id=data["list_id"],
        )
        return response_list
    else:
        e = ServiceInfoError(
            "Something went wrong. Please check provided Klaviyo keys."
        )
        if task_handle:
            task_handle.send_immediate_error(e)
        logger.exception(e)
        raise e