import lunardate
import datetime
from pqai_agent.logging_service import logger
from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool
from collections import defaultdict


class LunarFestivalMapper(BaseToolkit):
    # 常见农历节日定义（月份, 日期）
    FESTIVALS = {
        (1, 1): "春节",
        (1, 15): "元宵节",
        (2, 2): "龙抬头",
        (5, 5): "端午节",
        (7, 7): "七夕",
        (7, 15): "中元节",
        (8, 15): "中秋节",
        (9, 9): "重阳节",
        (12, 8): "腊八节",
        (12, 23): "小年",
        (12, 30): "除夕"
    }

    def __init__(self, year=2025):
        super().__init__()
        self.year = year
        self.festival_dates = self._calculate_festivals()

    def _calculate_festivals(self):
        """计算指定年份的农历节日对应的公历日期"""
        results = defaultdict(list)

        # 遍历整年的每一天
        start_date = datetime.date(self.year, 1, 1)
        end_date = datetime.date(self.year, 12, 31)
        current_date = start_date

        while current_date <= end_date:
            try:
                # 将公历转换为农历
                lunar = lunardate.LunarDate.fromSolarDate(
                    current_date.year,
                    current_date.month,
                    current_date.day
                )
                # 检查是否为农历节日（非闰月）
                festival_key = (lunar.month, lunar.day)
                if festival_key in self.FESTIVALS:
                    festival_name = self.FESTIVALS[festival_key]
                    results[festival_name].append(current_date)

            except ValueError:
                # 跳过无效日期（如2月30日等）
                pass

            # 下一天
            current_date += datetime.timedelta(days=1)

        # 处理结果（每个节日只取第一个出现的日期）
        return {name: dates[0] for name, dates in results.items()}

    def get_festival_date(self, festival_name):
        """获取指定节日的公历日期"""
        return self.festival_dates.get(festival_name, "节日未找到或不在该年")

    def get_all_festivals(self):
        """获取该年所有农历节日日期"""
        return self.festival_dates

    def get_tools(self):
        return [FunctionTool(self.get_festival_date)]