"""Data module utility functions.

集中管理数据模块中常用的工具函数，避免重复定义。
"""

from datetime import datetime, timedelta
from typing import Optional, List


# 默认全量同步开始日期
DEFAULT_START_DATE = "20180101"

# 今日日期 (YYYYMMDD 格式)
TODAY: str = datetime.now().strftime("%Y%m%d")


def get_today_date(cutoff_hour: int = 9) -> str:
    """获取今日日期（YYYYMMDD 格式）。

    考虑数据生成时间的逻辑：在 cutoff_hour 点之前，返回前一天的日期，
    因为当天的数据还未生成。A股数据通常在交易日收盘后（约 15:00-19:00）
    生成，但为了保险起见，默认使用早上 9 点作为分界。

    Args:
        cutoff_hour: 时间分界点（小时，24小时制），默认为 9。
                      当前时间小于此值时，返回前一天日期。

    Returns:
        日期字符串，格式为 YYYYMMDD

    Example:
        >>> # 假设当前是 2024-01-15 08:30
        >>> get_today_date()  # 返回 '20240114'（前一天）
        >>>
        >>> # 假设当前是 2024-01-15 10:00
        >>> get_today_date()  # 返回 '20240115'（当天）
        >>>
        >>> # 使用自定义分界点
        >>> get_today_date(cutoff_hour=15)  # 15点前返回前一天
    """
    now = datetime.now()
    if now.hour < cutoff_hour:
        # 在分界点之前，返回前一天
        prev_dt = now - timedelta(days=1)
        return prev_dt.strftime("%Y%m%d")
    return now.strftime("%Y%m%d")


def get_next_date(date_str: str) -> str:
    """获取给定日期的下一天。

    Args:
        date_str: YYYYMMDD 格式的日期

    Returns:
        YYYYMMDD 格式的下一天日期
    """
    dt = datetime.strptime(date_str, "%Y%m%d")
    next_dt = dt + timedelta(days=1)
    return next_dt.strftime("%Y%m%d")


def get_prev_date(date_str: str) -> str:
    """获取给定日期的前一天。

    Args:
        date_str: YYYYMMDD 格式的日期

    Returns:
        YYYYMMDD 格式的前一天日期
    """
    dt = datetime.strptime(date_str, "%Y%m%d")
    prev_dt = dt - timedelta(days=1)
    return prev_dt.strftime("%Y%m%d")


def parse_date(date_str: str) -> datetime:
    """解析 YYYYMMDD 格式的日期字符串。

    Args:
        date_str: YYYYMMDD 格式的日期

    Returns:
        datetime 对象
    """
    return datetime.strptime(date_str, "%Y%m%d")


def format_date(dt: datetime) -> str:
    """将 datetime 对象格式化为 YYYYMMDD 字符串。

    Args:
        dt: datetime 对象

    Returns:
        YYYYMMDD 格式的日期字符串
    """
    return dt.strftime("%Y%m%d")


def is_quarter_end(date_str: str) -> bool:
    """判断是否为季度最后一天。

    Args:
        date_str: YYYYMMDD 格式的日期

    Returns:
        是否为季度最后一天
    """
    month_day = date_str[4:]
    return month_day in ("0331", "0630", "0930", "1231")


def date_to_quarter(date_str: str) -> str:
    """将日期转换为对应季度的最后一天。

    Args:
        date_str: YYYYMMDD 格式的日期

    Returns:
        季度最后一天，格式为 YYYYMMDD
        例如: 20230115 -> 20230331
    """
    year = date_str[:4]
    month = int(date_str[4:6])

    if month <= 3:
        return year + "0331"
    elif month <= 6:
        return year + "0630"
    elif month <= 9:
        return year + "0930"
    else:
        return year + "1231"


def get_quarters_in_range(start_date: str, end_date: str) -> List[str]:
    """获取日期范围内的所有季度列表。

    Args:
        start_date: 开始日期 YYYYMMDD
        end_date: 结束日期 YYYYMMDD

    Returns:
        季度列表，格式为 YYYYMMDD，按时间倒序排列
        例如: ['20231231', '20230930', '20230630', '20230331']
    """
    quarters = []

    # 将开始日期和结束日期都转换为季度
    start_quarter = date_to_quarter(start_date)
    end_quarter = date_to_quarter(end_date)

    # 解析年份
    start_year = int(start_date[:4])
    end_year = int(end_date[:4])

    # 遍历所有年份和季度
    for year in range(end_year, start_year - 1, -1):
        year_str = str(year)
        # 季度顺序: Q4, Q3, Q2, Q1 (倒序)
        for quarter in ["1231", "0930", "0630", "0331"]:
            quarter_date = year_str + quarter

            # 只包含在范围内的季度
            if quarter_date >= start_quarter and quarter_date <= end_quarter:
                quarters.append(quarter_date)

    return quarters
