from airflow import DAG
from airflow.providers.standard.operators.python import PythonOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime
import sys

sys.path.append("/opt/airflow/src")

from config.assets import ASSETS

START_DATE = "2015-01-01"


def run_ingestion(asset):
    from ingestion.download_data import download_data, save_to_parquet
    df = download_data(asset, START_DATE)
    save_to_parquet(df, asset)


def run_cleaning(asset):
    from processing.clean_data import clean_data
    clean_data(asset)


def run_features(asset):
    from processing.feature_engineering import create_features
    create_features(asset)


def run_regime(asset):
    from models.hmm import classify_market_regime
    classify_market_regime(asset)


def notify_task(asset):
    from telegram.notify_regime import notify_today_regime
    notify_today_regime(asset)


with DAG(
    dag_id="market_regime_pipeline",
    start_date=datetime(2026, 1, 1),
    schedule="@daily",
    catchup=False,
    tags=["market", "pipeline"],
) as dag:

    for asset in ASSETS:

        with TaskGroup(group_id=f"{asset.replace('-', '_')}") as asset_group:

            t1_ingest = PythonOperator(
                task_id="ingestion",
                python_callable=run_ingestion,
                op_args=[asset],
            )

            t2_clean = PythonOperator(
                task_id="cleaning",
                python_callable=run_cleaning,
                op_args=[asset],
            )

            t3_features = PythonOperator(
                task_id="feature_engineering",
                python_callable=run_features,
                op_args=[asset],
            )

            t4_regime = PythonOperator(
                task_id="market_regime_classification",
                python_callable=run_regime,
                op_args=[asset],
            )

            t5_notify = PythonOperator(
                task_id="notify_regime",
                python_callable=notify_task,
                op_args=[asset],
            )

            t1_ingest >> t2_clean >> t3_features >> t4_regime >> t5_notify