from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.sensors.filesystem import FileSensor
from airflow.decorators import task, task_group

with DAG(
    'cafe_monthly_full_pipeline',
    default_args=default_args,
    description='서울 카페 월간 파이프라인 (Streaming & GPU Optimized)',
    schedule_interval='0 3 1 * *',  # 매월 1일 03:00 (KST)
    start_date=datetime(2026, 1, 25),
    catchup=False,
    tags=['cafe', 'monthly', 'streaming', 'gpu_pool'],
    max_active_runs=1,
) as dag:

with DAG(
    'cafe_weekly_review_pipeline',
    default_args=default_args,
    description='서울 카페 주간 리뷰 업데이트 파이프라인 (text-only)',
    schedule_interval='0 2 * * 6',  # 매주 토요일 02:00 (KST)
    start_date=datetime(2026, 1, 25),
    catchup=False,
    tags=['cafe', 'weekly', 'review', 'text-only'],
    max_active_runs=1,
) as dag:

    # 리뷰만 크롤링 (이미지 스킵)
    crawl_reviews = BashOperator(
        task_id='crawl_reviews',
        bash_command=f'''
            cd{CRAWLER_PATH} &&\
            python main.py --mode parallel --target text-only
        ''',
    )

    # 리뷰 요약
    summarize = BashOperator(
        task_id='summarize_reviews',
        bash_command=f'cd{PIPELINE_PATH} && python _03_summarize_reviews.py',
    )

    # DB 업데이트
    seed_reviews = BashOperator(
        task_id='seed_reviews',
        bash_command=f'cd{PROJECT_PATH} && python backend/core/seed_weekly_reviews.py',
    )

    # 의존성 정의
    crawl_reviews >> summarize >> seed_reviews

# 1. 처리할 동네 목록 동적 로드
@task
def get_target_dongs():
    dong_file = f"{CRAWLER_PATH}/seoul_parts/seoul_dongs.json"
    with open(dong_file, 'r', encoding='utf-8') as f:
        dongs = json.load(f)
    return sorted(dongs)  # ['강남구_개포1동', '강남구_개포2동', ...]

# 2. 동네별 Task Group 정의
@task_group(group_id='process_dong')
def process_neighborhood(dong_name):
    # 각 동네별 5단계 파이프라인
    step0 >> step1 >> step2 >> step3 >> step4

# 3. Dynamic Mapping으로 병렬 확장
target_dongs = get_target_dongs()
processed_groups = process_neighborhood.expand(dong_name=target_dongs)

@task_group(group_id='process_dong')
def process_neighborhood(dong_name):

    # A. 파일 감지 - 해당 동네 JSON이 생성되면 다음 단계 시작
    wait_for_json = FileSensor(
        task_id='wait_for_json',
        filepath=f"{RAW_JSON_DIR}/{dong_name}.json",
        fs_conn_id='fs_default',
        mode='reschedule',        # 대기 중에는 Worker 슬롯 반환
        poke_interval=60 * 5,     # 5분마다 파일 존재 확인
        timeout=60 * 60 * 24,     # 최대 24시간 대기
    )

    # B. 파일 감지 후 처리 시작
    step0_download = BashOperator(...)

    wait_for_json >> step0_download >> step1 >> step2 >> step3 >> step4

# GPU 작업은 별도 풀로 제한
step1_clip = BashOperator(
    task_id='step1_clip',
    bash_command=params['cmd_step1'],
    pool='gpu_pool',  # GPU 풀 (슬롯 제한)
)

step2_vlm = BashOperator(
    task_id='step2_vlm',
    bash_command=params['cmd_step2'],
    pool='gpu_pool',  # GPU 풀
)

step3_summary = BashOperator(
    task_id='step3_summary',
    bash_command=params['cmd_step3'],
    pool='gpu_pool',  # GPU 풀
)

step4_embed = BashOperator(
    task_id='step4_embed',
    bash_command=params['cmd_step4'],
    pool='gpu_pool',  # GPU 풀
)

@task_group(group_id='process_dong')
def process_neighborhood(dong_name):

    # 파라미터 준비
    @task
    def prepare_params(name):
        return {
            'sensor_path': f"{RAW_JSON_DIR}/{str(name)}.json",
            'cmd_step0': f'cd{PIPELINE_PATH} && python _00_image_extraction.py --dong-name{str(name)}',
            'cmd_step1': f'cd{PIPELINE_PATH} && python _01_extract_representative_images.py --dong-name{str(name)}',
            'cmd_step2': f'cd{PIPELINE_PATH} && python _02_generate_vlm_descriptions.py --dong-name{str(name)}',
            'cmd_step3': f'cd{PIPELINE_PATH} && python _03_summarize_reviews.py --dong-name{str(name)}',
            'cmd_step4': f'cd{PIPELINE_PATH} && python _04_generate_embeddings.py --dong-name{str(name)}',
        }

    params = prepare_params(dong_name)

    # 파일 감지
    wait_for_json = FileSensor(...)

    # 이미지 다운로드 (Default Pool)
    step0_download = BashOperator(
        task_id='step0_download',
        bash_command=params['cmd_step0'],
        pool='default_pool',
    )

    # GPU 작업 체인 (GPU Pool)
    step1_clip = BashOperator(task_id='step1_clip', pool='gpu_pool', ...)
    step2_vlm = BashOperator(task_id='step2_vlm', pool='gpu_pool', ...)
    step3_summary = BashOperator(task_id='step3_summary', pool='gpu_pool', ...)
    step4_embed = BashOperator(task_id='step4_embed', pool='gpu_pool', ...)

    # 의존성 체인
    wait_for_json >> step0_download >> step1_clip >> step2_vlm >> step3_summary >> step4_embed

# 실행 흐름
target_dongs = get_target_dongs()
processed_groups = process_neighborhood.expand(dong_name=target_dongs)

# 크롤링과 병렬로 진행, 모든 처리 완료 후 DB 적재
crawl_full >> seed_db
processed_groups >> seed_db