import uuid
from time import sleep

import requests

from service.conf import settings
from service.utils.aws_connectors import run_query


class AirflowEnrichmentRunner:
    def __init__(
        self, user_id, collection_id, workspace_schema=None, alliance_schema=None
    ):
        self.user_id = user_id
        self.collection_id = collection_id
        self.workspace_schema = workspace_schema
        self.alliance_schema = alliance_schema
        self.schema = self.workspace_schema or self.alliance_schema
        if self.schema is None:
            raise ValueError(
                'One of ["workspace_schema", "alliance_schema"] cannot be None.'
            )

    def _make_airflow_api_request(self, url, method, payload=None):
        response = requests.request(
            method=method,
            url=settings.AIRFLOW_CREDENTIALS["host"] + url,
            json=payload,
            auth=(
                settings.AIRFLOW_CREDENTIALS["username"],
                settings.AIRFLOW_CREDENTIALS["password"],
            ),
        )
        if response.status_code != 200:
            raise RuntimeError(response.text)
        return response

    def execute_collection_enrichment(self):
        if self.workspace_schema is None:
            raise ValueError(
                "workspace_schema parameter cannot be None to run enrichments_dag"
            )

        payload = {
            "conf": {
                "profile": settings.APPSYNC_PROFILE,
                "workspace_schema": self.workspace_schema,
                "user_id": self.user_id,
                "collection_id": self.collection_id,
            }
        }

        self._execute_enrichment(dag_id="enrichments_dag", payload=payload)

    def execute_alliance_enrichment(self):
        if self.alliance_schema is None:
            raise ValueError(
                "alliance_schema parameter cannot be None to run algorithms_dag"
            )

        payload = {
            "conf": {
                "alliance_schema": self.alliance_schema,
                "profile": settings.APPSYNC_PROFILE,
                "user_id": self.user_id,
                "collection_id": self.collection_id,
            }
        }

        self._execute_enrichment(dag_id="algorithms_dag", payload=payload)

    def execute_segment_enrichment(self):
        payload = {
            "conf": {
                "profile": settings.APPSYNC_PROFILE,
                "user_id": self.user_id,
                "collection_id": self.collection_id,
            }
        }
        if self.workspace_schema is not None:
            payload["conf"].update({"workspace_schema": self.workspace_schema})
        elif self.alliance_schema is not None:
            payload["conf"].update({"alliance_schema": self.alliance_schema})

        self._execute_enrichment(dag_id="algorithms_dag", payload=payload)

    def _log_enrichment_run(self, dag_id, dag_run_id, execution_date):
        run_query(
            f"""
                INSERT INTO {self.schema}.collection_airflow
                VALUES (%(collection_id)s, %(dag_id)s, %(dag_run_id)s, %(execution_date)s,
                        (SELECT COUNT(cf.fan_id)
                         FROM {self.schema}.collection_fan cf
                         WHERE cf.collection_id = %(collection_id)s)
                        )
            """,
            dict(
                collection_id=self.collection_id,
                dag_id=dag_id,
                dag_run_id=dag_run_id,
                execution_date=execution_date,
            ),
        )

    def _get_fans_count(self):
        [[count]] = run_query(
            f"""
                SELECT COUNT(cf.fan_id)
                FROM {self.schema}.collection_fan cf
                WHERE cf.collection_id = %(collection_id)s
            """,
            dict(collection_id=self.collection_id),
        )
        return count

    def _get_source_collection_attributes(self):
        if self.workspace_schema is not None:
            query_result = run_query(
                f"""
                    SELECT ARRAY_AGG(DISTINCT ca.attribute_id)
                    FROM {self.schema}.collection_attribute ca
                         JOIN {self.schema}.attribute a ON a.id = ca.attribute_id
                    WHERE (ca.collection_id = %(collection_id)s  -- if collection is a source collection
                            OR ca.collection_id IN (SELECT collection_id -- if collection is a segment / superset
                                                    FROM {self.schema}.set_collection
                                                    WHERE set_id = %(collection_id)s))
                      AND ca.attribute_id < 10000
                      AND a.name NOT LIKE 'enr%%'
                """,
                dict(collection_id=self.collection_id),
            )
        else:  # self.alliance_schema is not None
            query_result = run_query(
                f"""
                    SELECT ARRAY_AGG(DISTINCT ca.attribute_id)
                    FROM {self.schema}.collection_attribute ca
                         JOIN {self.schema}.attribute a ON a.id = ca.attribute_id
                    WHERE ca.collection_id IN (SELECT c.id
                                               FROM {self.schema}.collection c
                                               WHERE collection_type = 'source')
                      AND ca.attribute_id < 10000
                      AND a.name NOT LIKE 'enr%%'
                """
            )
        if len(query_result) > 0:
            [[attribute_ids]] = query_result
        else:
            attribute_ids = []
        return attribute_ids

    def _execute_enrichment(self, dag_id, payload):
        payload["dag_run_id"] = f"{uuid.uuid4()}:{self.schema}:{self._get_fans_count()}"
        payload["conf"][
            "source_attribute_ids"
        ] = self._get_source_collection_attributes()
        api_response = self._make_airflow_api_request(
            url=f"/dags/{dag_id}/dagRuns", method="POST", payload=payload
        ).json()

        self._wait_for_enrichment_initialization_complete(
            api_response["dag_id"], api_response["dag_run_id"]
        )

        self._log_enrichment_run(
            api_response["dag_id"],
            api_response["dag_run_id"],
            api_response["execution_date"],
        )

    def _wait_for_enrichment_initialization_complete(
        self, dag_id, dag_run_id, delay=1, max_retries=10
    ):
        """
        Wait until Airflow creates task_instances, which are necessary to estimate
        enrichment runtime for the progress bar. Looks like the scheduler starts some
        tasks before communicating the details about them to Airflow.
        """
        retries = 0
        while True:
            api_response = self._make_airflow_api_request(
                url=f"/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
                method="GET",
            ).json()

            if len(api_response.get("task_instances", [])) > 0:
                return True

            if retries >= max_retries:
                raise RuntimeError(f"DAGRun {dag_run_id} has failed to initialize.")

            retries += 1
            sleep(delay)