import importlib
from pathlib import PurePosixPath

from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.utils.task_group import TaskGroup

from common.config.app import ELASTIC_SEARCH_URL
from common.config.search import ALL_DSPS
from common.operators.switch_es import SwitchElasticsearchDataSource
from common.operators.transform_search import TransformSearchOperator
from common.transfers.s3_to_es import S3ToElasticSearchOperator
from common.transfers.s3_to_pg import S3ToPostgresOperator


__all__ = ["SearchTaskManager"]


class SearchTaskManager:

    def __init__(self, source: str, transform_pool: str = "transform_pool", transform_pool_slots: int = 5):
        self.source = source

        self.transform_pool = transform_pool
        self.transform_pool_slots = transform_pool_slots

        # TODO: change suffix, remove prefix
        self.table_name = f"test_{source}_search_data_airflow"
        self.table_name_alt = f"{self.table_name}_alt"
        self.base_index_name = "test_{{ params.config.app.ELASTIC_SEARCH_ENV }}_{{ params.source }}_search_airflow"
        self.index_name = self.base_index_name + "_{{ get_timestamp(dag_run) }}"

        self.base_path = PurePosixPath("search", source)
        self.raw_data_path = self.base_path / "raw_data"
        self.agg_data_path = self.base_path / "agg_data"
        self.index_data_path = self.base_path / "index_data"

    def get_process_task(self):
        source_config = importlib.import_module(f"common.config.search.{self.source}")

        with TaskGroup(self.source) as process:
            with TaskGroup("extract") as extract:
                for dsp in ALL_DSPS:
                    SnowflakeOperator(
                        task_id=dsp,
                        session_parameters={"timezone": "UTC"},
                        sql=f"{self.base_path}/{dsp}/extract.sql",
                        params={
                            "raw_data_path": self.raw_data_path / dsp
                        },
                    )

            recreate_table = PostgresOperator(
                task_id="recreate_table",
                sql="postgres/recreate_table.sql",
                params={
                    "dst_table": self.table_name_alt
                }
            )

            transform = TransformSearchOperator(
                task_id="transform",
                pool=self.transform_pool,
                pool_slots=self.transform_pool_slots,
                raw_data_path=self.raw_data_path,
                agg_data_path=self.agg_data_path,
                index_data_path=self.index_data_path,
                top_props=source_config.TOP_PROPS,
                index_props=source_config.INDEX_PROPS,
                dsp_index_props_map=source_config.DSP_INDEX_PROPS_MAP,
                dsp_to_interpolate=source_config.DSP_TO_INTERPOLATE,
                dsp_to_extend=source_config.DSP_TO_EXTEND,
                dsp_change_trends_props_map=source_config.DSP_CHANGE_TRENDS_PROPS_MAP
            )

            with TaskGroup("load") as load:
                S3ToPostgresOperator(
                    task_id="data",
                    upload_sql="transfers/s3_to_pg.sql",
                    recreate_sql="postgres/recreate_table.sql",
                    params={
                        "dst_table": self.table_name_alt
                    },
                    s3_data_path=self.index_data_path
                )
                S3ToElasticSearchOperator(
                    task_id="index",
                    s3_data_path=self.index_data_path,
                    elastic_search_hosts=[ELASTIC_SEARCH_URL],
                    elastic_search_index_name=self.index_name,
                    elastic_search_index_mapping=source_config.INDEX_MAPPING,
                )

            extract >> recreate_table >> transform >> load

        return process

    def get_switch_task(self):
        with TaskGroup(f"search_{self.source}") as switch:
            PostgresOperator(
                task_id="data",
                sql="postgres/switch_data_source.sql",
                params={
                    "src_table": self.table_name_alt,
                    "dst_table": self.table_name
                }
            )

            SwitchElasticsearchDataSource(
                task_id="index",
                hosts=[ELASTIC_SEARCH_URL],
                index_name=self.index_name,
                base_index_name=self.base_index_name
            )

        return switch