from datetime import datetime

import pandas as pd
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.postgres.hooks.postgres import PostgresHook

# Список баз данных PostgreSQL, из которых извлекаются данные.
DATABASES = ['postgres_db1', 'postgres_db2']

# Название таблицы, из которой извлекаются данные.
TABLE_NAME = 'vinoteka'

# Название S3-бакета, куда загружаются объединенные данные.
BUCKET = 'postgres-data-csv'


def extract_data(ds, **kwargs):
    """
    Извлекает данные из указанных баз данных PostgreSQL и объединяет их в один DataFrame.

    Args:
        ds: Шаблон даты (например, {{ ds }}), передаваемый Airflow.

    Возвращаемое значение:
        Объединенные данные сохраняются в виде pickle-файла '/tmp/combined_data.pkl'.
    """
    databases = DATABASES
    dataframes = []

    # Цикл по списку баз данных и извлечение данных из каждой.
    for db in databases:
        hook = PostgresHook(postgres_conn_id=db)
        sql = f"SELECT * FROM {TABLE_NAME}"
        df = hook.get_pandas_df(sql)
        dataframes.append(df)
    
    # Объединение данных из всех баз данных в один DataFrame.
    combined_data = pd.concat(dataframes)
    
    # Сохранение объединенных данных в локальный файл.
    combined_data.to_pickle('/tmp/combined_data.pkl')


def combine_data(ds, **kwargs):
    """
    Дополнительно обрабатывает объединенные данные (если необходимо) и сохраняет их.

    Args:
        ds: Шаблон даты (например, {{ ds }}), передаваемый Airflow.

    Возвращаемое значение:
        Обработанные данные сохраняются в виде pickle-файла '/tmp/merged_data.pkl'.
    """
    combined_data = pd.read_pickle('/tmp/combined_data.pkl')
    
    # Дополнительные операции по обработке данных можно добавить здесь.
    
    # Сохранение обработанных данных в локальный файл.
    combined_data.to_pickle('/tmp/merged_data.pkl')


def upload_to_minio(ds, **kwargs):
    """
    Загружает обработанные данные в S3-совместимый MinIO хранилище в виде CSV-файла.

    Args:
        ds: Шаблон даты (например, {{ ds }}), передаваемый Airflow.

    Возвращаемое значение:
        CSV-файл загружается в указанный S3-бакет с уникальным именем.
    """
    merged_data = pd.read_pickle('/tmp/merged_data.pkl')
    
    # Определение пути к локальному CSV-файлу.
    csv_path = '/tmp/merged_data.csv'
    
    # Сохранение данных в CSV-файл.
    merged_data.to_csv(csv_path, index=False)
    
    # Создание уникального имени для загружаемого файла, основанного на текущем времени.
    now = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    object_name = f'merged_data_{now}.csv'
    
    # Подключение к MinIO и загрузка файла в S3-бакет.
    minio_hook = S3Hook('minio_conn_id')
    minio_hook.load_file(
        bucket_name=BUCKET,
        key=object_name,
        filename=csv_path,
        replace=True,
    )

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.utils.dates import days_ago

from .functions import extract_data, combine_data, upload_to_minio

# Определение аргументов по умолчанию для всех задач в DAG
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': days_ago(1),
    'email_on_failure': False,
    'email_on_retry': False,
}

# Создание DAG для извлечения данных, их объединения и загрузки в MinIO
with DAG(
    'data_collection_processing_and_writing_to_csv',
    default_args=default_args,
    description='Extract data, combine and upload to minio',
    schedule_interval='@daily',
) as dag:
    
    # Оператор для извлечения данных из баз данных PostgreSQL
    extract_data_task = PythonOperator(
        task_id='extract_data',
        python_callable=extract_data,
        provide_context=True,
    )
    
    # Оператор для объединения и обработки данных
    combine_data_task = PythonOperator(
        task_id='combine_data',
        python_callable=combine_data,
        provide_context=True,
    )

    # Оператор для загрузки обработанных данных в MinIO в формате CSV
    upload_to_minio_task = PythonOperator(
        task_id='upload_to_minio',
        python_callable=upload_to_minio,
        provide_context=True,
    )
    
    # Определение порядка выполнения задач
    extract_data_task >> combine_data_task >> upload_to_minio_task