# Импортируем необходимые компоненты из airflow
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

# Импортируем библиотеки для работы с разными источниками данных
import psycopg2                    # Для PostgreSQL
import boto3                       # Для подключения к Amazon S3
import fsspec                      # Универсальный доступ к файловым системам, включая NFS
from clickhouse_driver import Client  # Для подключения к ClickHouse
import os                          # Работа с файловой системой

# Функция для получения данных из PostgreSQL
def fetch_postgresql(**kwargs):
    # Устанавливаем соединение с Postgres
    conn = psycopg2.connect(
        dbname='your_db', user='user', password='pass',
        host='localhost', port=5432
    )
    # Создаём курсор для выполнения SQL-запроса
    cur = conn.cursor()
    # Выполняем SQL-запрос
    cur.execute("SELECT * FROM your_table LIMIT 10;")
    # Забираем результаты
    result = cur.fetchall()
    # Закрываем соединение
    cur.close()
    conn.close()
    # Возвращаем результат
    return result

# Функция для получения данных из ClickHouse
def fetch_clickhouse(**kwargs):
    # Подключаемся к ClickHouse-серверу
    client = Client('localhost')
    # Выполняем SQL-запрос
    data = client.execute('SELECT * FROM your_table LIMIT 10')
    # Возвращаем результат
    return data

# Функция для получения данных из AWS S3
def fetch_s3(**kwargs):
    # Создаём клиент автомата S3
    s3 = boto3.client('s3')
    # Загружаем объект (например, JSON-файл)
    obj = s3.get_object(Bucket='your-bucket', Key='example.json')
    # Считываем содержимое файла
    data = obj['Body'].read()
    # Возвращаем данные (в байтах)
    return data

# Функция для чтения файла с NFS-ресурса
def fetch_nfs_file(**kwargs):
    # Путь к файлу на подключённом сетевом хранилище
    file_path = '/mnt/nfs/data_file.txt'
    # Читаем содержимое файла
    with open(file_path, 'r') as f:
        return f.read()

# Функция для чтения локального файла
def fetch_local_file(**kwargs):
    # Путь к файлу в локальной ФС
    file_path = 'local_data.txt'
    # Читаем содержимое
    with open(file_path, 'r') as f:
        return f.read()

# Функция для агрегации данных из всех источников
def aggregate_and_respond(**context):
    # Вытаскиваем данные из предыдущих задач через XCom
    pg_data = context['ti'].xcom_pull(task_ids='fetch_postgresql')
    ch_data = context['ti'].xcom_pull(task_ids='fetch_clickhouse')
    s3_data = context['ti'].xcom_pull(task_ids='fetch_s3')
    nfs_data = context['ti'].xcom_pull(task_ids='fetch_nfs_file')
    file_data = context['ti'].xcom_pull(task_ids='fetch_local_file')
    
    # Возвращаем агрегированные данные в виде словаря
    return {
        "PostgreSQL": pg_data,
        "ClickHouse": ch_data,
        "S3": s3_data.decode() if isinstance(s3_data, bytes) else s3_data,  # Преобразуем байты в строку
        "NFS": nfs_data,
        "File": file_data,
    }

# Настройки DAG: владелец, дата запуска (в прошлом, чтобы не было catchup)
default_args = {
    'start_date': datetime(2023, 1, 1),
    'owner': 'ai_assistant',
}

# Определяем DAG — главный объект пайплайна
with DAG(
    'multi_source_dag_pipeline',         # Уникальное имя пайплайна
    default_args=default_args,           # Параметры по умолчанию
    schedule_interval=None,              # Запускается вручную
    catchup=False                        # Не выполняет предыдущие пропущенные запуски
) as dag:
    
    # Задача: получить данные из PostgreSQL
    fetch_postgresql_task = PythonOperator(
        task_id='fetch_postgresql',        # ID задачи
        python_callable=fetch_postgresql   # Указание, какую функцию вызывать
    )

    # Задача: получить данные из ClickHouse
    fetch_clickhouse_task = PythonOperator(
        task_id='fetch_clickhouse',
        python_callable=fetch_clickhouse
    )