import logging

import airflow
from airflow.exceptions import AirflowSkipException
from airflow.models.baseoperator import BaseOperator

from fansifter.constants import ALGORITHMS
from fansifter.operators.base import FanSifterLambdaOperator
from fansifter.hooks.collection import CollectionHook


class EnrichmentOperator(FanSifterLambdaOperator):
    def __init__(self, enrichment_name, enrichment_config=None, output_xcom_key_suffix="enrichment", **kwargs) -> None:
        super().__init__(**kwargs)
        self.enrichment_name = enrichment_name
        self.enrichment_config = enrichment_config if enrichment_config is not None else {}
        self.output_xcom_key_suffix = output_xcom_key_suffix

    def execute(self, context):
        ti: airflow.models.TaskInstance = context["task_instance"]

        conf = self._get_config()

        profile = conf["profile"]
        workspace_schema = conf.get("workspace_schema")
        alliance_schema = conf.get("alliance_schema")
        user_id = conf["user_id"]
        collection_id = conf["collection_id"]

        payload = {
            "endpoint": f"enrich/{self.enrichment_name}",
            "arguments": {
                "user_id": user_id,
                "collection_id": collection_id,
                "enrichment_config": self.enrichment_config,
            },
        }
        if workspace_schema is not None:
            payload["arguments"].update({"workspace_schema": workspace_schema})
        elif alliance_schema is not None:
            payload["arguments"].update({"alliance_schema": alliance_schema})

        try:
            lambda_result = self.invoke_lambda(
                function_name=f"frontend-api-{profile}-enrichment-requests", payload=payload
            )
            if lambda_result is None:
                logging.info("Enrichment has not created any data, terminating the task...")
                return None
        except Exception as e:
            logging.error(e)
            raise e

        ti.xcom_push(key=f"{self.enrichment_name}-{self.output_xcom_key_suffix}", value=lambda_result)


class SaveEnrichmentDataOperator(FanSifterLambdaOperator):
    def __init__(self, enrichment_name, input_task_id, input_xcom_key_suffix="enrichment", **kwargs) -> None:
        super().__init__(**kwargs)
        self.enrichment_name = enrichment_name
        self.input_task_id = input_task_id
        self.input_xcom_key_suffix = input_xcom_key_suffix

    def execute(self, context):
        ti: airflow.models.TaskInstance = context["task_instance"]

        conf = self._get_config()
        profile = conf["profile"]
        schema = conf.get("workspace_schema") or conf.get("alliance_schema")

        [enrichment_result] = ti.xcom_pull(
            key=f"{self.enrichment_name}-{self.input_xcom_key_suffix}", task_ids=[self.input_task_id]
        )
        if enrichment_result is None or enrichment_result.get("csv") is None:
            logging.info("No data to save, skipping the task...")
            raise AirflowSkipException

        target_collection_id = enrichment_result.get("collection_id")
        file_path = enrichment_result.get("csv")

        payload = {
            "schema": schema,
            "table": "fan_attribute",
            "collection_id": target_collection_id,
            "set_finished_status": self.enrichment_name not in ALGORITHMS,
            "file_path": file_path,
        }

        try:
            self.invoke_lambda(function_name=f"frontend-api-{profile}-save-fan-attribute-data", payload=payload)
        except Exception as e:
            logging.error(e)
            raise e


class CreateEnrichmentSegmentsOperator(FanSifterLambdaOperator):
    def __init__(self, enrichment_name, input_task_id, input_xcom_key_suffix="enrichment", **kwargs) -> None:
        super().__init__(**kwargs)
        self.enrichment_name = enrichment_name
        self.input_task_id = input_task_id
        self.input_xcom_key_suffix = input_xcom_key_suffix

    def execute(self, context):
        ti: airflow.models.TaskInstance = context["task_instance"]

        conf = self._get_config()
        profile = conf["profile"]
        workspace_schema = conf.get("workspace_schema")
        alliance_schema = conf.get("alliance_schema")
        user_id = conf["user_id"]
        collection_id = conf["collection_id"]

        [enrichment_result] = ti.xcom_pull(
            key=f"{self.enrichment_name}-{self.input_xcom_key_suffix}", task_ids=[self.input_task_id]
        )
        if enrichment_result is None or enrichment_result.get("csv") is None:
            logging.info("Enrichment has not created any data, skipping the task...")
            raise AirflowSkipException

        segment_attribute = enrichment_result.get("segment_attribute")
        segment_name = enrichment_result.get("segment_name")
        segment_values = enrichment_result.get("segment_values")
        algo_collection_id = enrichment_result.get("collection_id")

        payload = {
            "endpoint": "post_enrich/create_segments",
            "arguments": {
                "user_id": user_id,
                "collection_id": collection_id,
                "algo_collection_id": algo_collection_id,
                "segment_attribute": segment_attribute,
                "segment_name": segment_name,
                "segment_values": segment_values,
            },
        }
        if workspace_schema is not None:
            payload["arguments"].update({"workspace_schema": workspace_schema})
        elif alliance_schema is not None:
            payload["arguments"].update({"alliance_schema": alliance_schema})

        try:
            self.invoke_lambda(function_name=f"frontend-api-{profile}-post-enrichment-requests", payload=payload)
        except Exception as e:
            logging.error(e)
            raise e


class ChangeCollectionStatusOperator(FanSifterLambdaOperator):
    def __init__(self, status, **kwargs) -> None:
        super().__init__(**kwargs)
        self.status = status

    def execute(self, context):
        conf = self._get_config()
        profile = conf["profile"]
        workspace_schema = conf.get("workspace_schema")
        alliance_schema = conf.get("alliance_schema")
        user_id = conf["user_id"]
        collection_id = conf["collection_id"]

        collection_hook = CollectionHook(
            collection_id=collection_id,
            user_id=user_id,
            profile=profile,
            workspace_schema=workspace_schema,
            alliance_schema=alliance_schema,
        )

        try:
            collection_hook.change_status(self.status)
        except Exception as e:
            logging.error(e)
            raise e