from abc import ABC, abstractmethod, abstractproperty

from airflow.models.dag import DAG

from gcp_airflow_foundations.base_class.data_source_table_config import (
    DataSourceTablesConfig,
)
from gcp_airflow_foundations.base_class.source_table_config import SourceTableConfig
from gcp_airflow_foundations.base_class.source_config import SourceConfig
from gcp_airflow_foundations.enums.schema_source_type import SchemaSourceType
from gcp_airflow_foundations.common.gcp.load_builder import load_builder
from gcp_airflow_foundations.source_class.schema_source_config import (
    AutoSchemaSourceConfig,
    GCSSchemaSourceConfig,
    BQLandingZoneSchemaSourceConfig,
)

import logging


class DagBuilder(ABC):
    """A base DAG builder for creating a list of DAGs for a given source.
    Attributes
    ----------
    sources : list(DagBuilder)
        List of initialized subclasses of this class (dynamically updated during runtime)
    config : DataSourceTablesConfig
        DAG configuration object
    """

    sources = []

    def __init__(self, default_task_args: dict, config: DataSourceTablesConfig):
        self.default_task_args = default_task_args
        self.config = config
        self.validate_extra_options()
        self.set_schema_method_type()
        self.validate_custom_method()
        super().__init__()

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if not hasattr(cls, "source_type"):
            raise AttributeError(
                f"{cls} source class is missing 'source_type' attribute"
            )
        cls.sources.append(cls)

    def build_dags(self):
        """
        Main DAG building method
        """
        data_source = self.config.source
        logging.info(f"Building DAG for {data_source.name}")

        dags = []
        for table_config in self.config.tables:
            table_default_task_args = self.default_task_args_for_table(
                self.config, table_config
            )
            logging.info(f"table_default_task_args {table_default_task_args}")

            kwargs = data_source.dag_args if data_source.dag_args else {}

            with DAG(
                dag_id=f"{data_source.name}.{table_config.table_name}",
                description=f"{data_source.name} to BigQuery load for {table_config.table_name}",
                schedule_interval=data_source.ingest_schedule,
                default_args=table_default_task_args,
                catchup=data_source.catchup,
                render_template_as_native_obj=True,
                **kwargs,
            ) as dag:

                load_to_bq_landing = self.get_bq_ingestion_task(dag, table_config)

                self.get_datastore_ingestion_task(
                    dag, load_to_bq_landing, data_source, table_config
                )

                dags.append(dag)

        extra_dags = self.get_extra_dags()
        if extra_dags is not None and not extra_dags == []:
            dags = dags + self.get_extra_dags()
        return dags

    @abstractmethod
    def get_bq_ingestion_task(self, table_config: SourceTableConfig):
        """Abstract method for the Airflow task that ingests data to the BigQuery staging table"""
        pass

    def get_datastore_ingestion_task(
        self,
        dag,
        preceding_task,
        data_source: SourceConfig,
        table_config: SourceTableConfig,
    ):
        """Method for the Airflow task group that upserts data to the ODS and HDS tables"""
        load_builder(
            data_source=data_source,
            table_config=table_config,
            schema_config=BQLandingZoneSchemaSourceConfig,
            preceding_task=preceding_task,
            dag=dag,
        )

    def get_extra_dags(self):
        # Override if non-table ingestion DAGs are needed for the flow: return them here
        return []

    @abstractmethod
    def set_schema_method_type(self):
        pass

    def get_schema_method_class(self):
        if self.schema_source_type == SchemaSourceType.GCS:
            return GCSSchemaSourceConfig

        elif self.schema_source_type == SchemaSourceType.AUTO:
            return AutoSchemaSourceConfig

        else:
            return

    def validate_custom_method(self):
        if self.schema_source_type == SchemaSourceType.CUSTOM:
            assert (
                self.get_schema_method_class() is not None
            ), "A custom method for schema parsing must be provided"

    @abstractmethod
    def validate_extra_options(self):
        pass

    def default_task_args_for_table(
        self, config: DataSourceTablesConfig, table_config: SourceTableConfig
    ):
        return {
            **self.default_task_args,
            "start_date": config.table_start_date(table_config),
        }