#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
时区工具模块
提供时区转换相关的工具函数
"""

from datetime import datetime
import pytz


def get_beijing_ds(context):
    """
    获取北京时区的ds日期
    
    Args:
        context: Airflow任务上下文
        
    Returns:
        str: 北京时区的日期字符串，格式：YYYY-MM-DD
    """
    # 获取DAG调度时间
    execution_date = context.get('execution_date')
    ds = context.get('ds')
    
    print(f"DAG调度时间 (UTC): {execution_date}")
    print(f"DAG执行日期ds (UTC): {ds}")
    
    # 将调度时间转换为北京时区
    beijing_tz = pytz.timezone('Asia/Shanghai')
    beijing_time = execution_date.astimezone(beijing_tz)
    
    # 获取北京时区的日期字符串
    ds_beijing = beijing_time.strftime('%Y-%m-%d')
    print(f"DAG调度时间 (北京时区): {beijing_time}")
    print(f"DAG执行日期ds (北京时区): {ds_beijing}")
    
    return ds_beijing


def convert_utc_to_beijing(utc_datetime):
    """
    将UTC时间转换为北京时区
    
    Args:
        utc_datetime: UTC时间对象
        
    Returns:
        datetime: 北京时区的时间对象
    """
    beijing_tz = pytz.timezone('Asia/Shanghai')
    return utc_datetime.astimezone(beijing_tz)


def get_beijing_date_str(utc_datetime):
    """
    获取北京时区的日期字符串
    
    Args:
        utc_datetime: UTC时间对象
        
    Returns:
        str: 北京时区的日期字符串，格式：YYYY-MM-DD
    """
    beijing_time = convert_utc_to_beijing(utc_datetime)
    return beijing_time.strftime('%Y-%m-%d')
