import logging
import os
from datetime import UTC, datetime
from enum import StrEnum
from http import HTTPStatus
from typing import TypedDict

import requests

# In Airflow 3.x there will be:
# from airflow.sdk import dag, task
# This will be deprecated in 3.x:
from airflow.decorators import dag, task
from airflow.models import Variable
from airflow.operators.python import get_current_context
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.providers.trino.hooks.trino import TrinoException, TrinoHook


logger = logging.getLogger(__name__)


TRINO_CONN_ID = 'trino'
APP_CONN_ID = 'ray_db'
# Переменная окружения для подключения к веб-сервису Ray:
#   RAY_WEB_CONN_STR — базовый URL API (например: http://ray-gw:8080/api/v1)
# Airflow Variable для токена:
#   ray_api_token    — Bearer-токен технического пользователя
RAY_API_TOKEN_VAR = 'ray_api_token'


class QueryResultStatus(StrEnum):
    SUCCESS = "success"
    ERROR = "error"


class QueryRunStatus(StrEnum):
    QUEUED = 'queued'
    RUNNING = 'running'
    SUCCESS = 'success'
    FAILED = 'failed'


class QueryExecutionStatus(StrEnum):
    POSITIVE = 'POSITIVE'
    NEGATIVE = 'NEGATIVE'


class PositiveQueryExecution(TypedDict):
    status: QueryExecutionStatus
    result: str


class NegativeQueryExecution(TypedDict):
    status: QueryExecutionStatus
    error_msg: str


@dag(
    dag_id="trino",
    schedule=None,
    start_date=datetime(2024, 1, 1, tzinfo=UTC),
    catchup=False,
    is_paused_upon_creation=False,
    tags=["trino"],
)
def trino_dag():
    @task
    def execute_trino_query() -> str:
        # Получение контекста запуска DAG:
        context = get_current_context()
        query_sql = context['dag_run'].conf['sql']
        query_run_id = context['dag_run'].conf['query_run_id']

        trino_hook = TrinoHook(trino_conn_id=TRINO_CONN_ID)

        try:
            result = trino_hook.get_records(query_sql)
            return PositiveQueryExecution(
                status=QueryExecutionStatus.POSITIVE,
                result=next(iter(result))[0],
            )
        except TrinoException as e:
            logger.warning(f'Error: execute_trino_query with query_run_id: "{query_run_id}".')
            logger.warning('Query SQL:')
            logger.warning(query_sql)
            logger.warning(e)
            return NegativeQueryExecution(
                status=QueryExecutionStatus.NEGATIVE,
                error_msg=str(e),
            )

    @task
    def process_results(query_result: PositiveQueryExecution | NegativeQueryExecution) -> None:
        # Получение контекста запуска DAG:
        context = get_current_context()
        expected_result = context['dag_run'].conf['expected_result']
        query_run_id = context['dag_run'].conf['query_run_id']

        pg_hook = PostgresHook(postgres_conn_id=APP_CONN_ID)

        connection = pg_hook.get_conn()
        cursor = connection.cursor()

        if query_result['status'] == QueryExecutionStatus.POSITIVE and query_result['result'] is not None:
            logger.info('Successfully completed test. Preparing for save.')
            logger.info(
                f"queryrunid: {query_run_id}. Execution result: {query_result['result']}. "
                f"Expected result: {expected_result}.",
            )
            cursor.execute(
                "update ray.query_run"
                " set"
                " result = (%(query_result)s = %(expected_result)s),"
                " return_value = json_build_object('value', %(query_result)s),"
                " result_status = %(query_result_status)s,"
                " query_run_status = %(query_run_status)s"
                " where query_run_id = %(query_run_id)s",
                {
                    'query_result': query_result['result'],
                    'expected_result': expected_result,
                    'query_run_id': query_run_id,
                    'query_result_status': QueryResultStatus.SUCCESS,
                    'query_run_status': QueryRunStatus.SUCCESS,
                },
            )
        else:
            error_msg = (
                'Выполнение SQL теста не удалось. '
                'Выборка из источника данных оказалась пустой. '
                'Пожалуйста, проверьте параметры запроса и убедитесь, что данные доступны для анализа.'
                if 'result' in query_result and query_result['result'] is not None
                else
                query_result['error_msg']
            )

            cursor.execute(
                "update ray.query_run"
                " set"
                " error_msg = %(error_msg)s,"
                " result_status = %(query_result_status)s,"
                " query_run_status = %(query_run_status)s"
                " where query_run_id = %(query_run_id)s",
                {
                    'error_msg': error_msg,
                    'query_run_id': query_run_id,
                    'query_result_status': QueryResultStatus.ERROR,
                    'query_run_status': QueryRunStatus.FAILED,
                },
            )

        connection.commit()

    @task
    def create_incident() -> None:
        """Создание инцидента по результатам выполнения теста.

        Вызывает эндпоинт POST /incidents/from_query_run веб-сервиса Ray,
        который самостоятельно определяет необходимость создания инцидента
        и формирует его на основе данных query_run.
        """
        context = get_current_context()
        query_run_id = context['dag_run'].conf['query_run_id']

        api_url = os.environ['RAY_WEB_CONN_STR']
        api_token = Variable.get(RAY_API_TOKEN_VAR)

        response = requests.post(
            url=f'{api_url}/incidents/from_query_run',
            json={'query_run_id': query_run_id},
            headers={
                'Authorization': f'Bearer {api_token}',
                'Content-Type': 'application/json',
            },
            timeout=30,
        )

        if response.status_code == HTTPStatus.CREATED:
            body = response.json()
            if body is None:
                logger.info('Инцидент не требуется для query_run_id: %s.', query_run_id)
            else:
                logger.info('Инцидент создан для query_run_id: %s. Incident id: %s.', query_run_id, body.get('id'))
        else:
            logger.error(
                'Ошибка при создании инцидента для query_run_id: %s. Status: %s. Response: %s',
                query_run_id, response.status_code, response.text,
            )
            response.raise_for_status()

    query_results = execute_trino_query()
    processed_data = process_results(query_results)
    incident_created = create_incident()

    query_results >> processed_data >> incident_created


dag_instance = trino_dag()
