DataOps-platform/app/core/data_parse/calendar.py

"""
黄历信息数据模型
基于 create_calendar_info.sql 中的DDL定义创建
"""

from datetime import date
from typing import Optional
from sqlalchemy import Column, Integer, Date, Text, String
from sqlalchemy.orm import Session
from sqlalchemy import create_engine, text
import json
import requests
from app import db
from .calendar_config import CALENDAR_API_CONFIG


class CalendarInfo(db.Model):
    """
    黄历信息表数据模型
    
    对应数据库表: public.calendar_info
    表注释: 黄历信息表
    """
    __tablename__ = 'calendar_info'
    __table_args__ = {'schema': 'public'}
    
    # 主键ID (serial primary key)
    id = db.Column(db.Integer, primary_key=True, autoincrement=True, comment='主键ID')
    
    # 阳历日期 (date not null)
    yangli = db.Column(db.Date, nullable=False, comment='阳历日期')
    
    # 阴历日期 (text not null)
    yinli = db.Column(db.Text, nullable=False, comment='阴历日期')
    
    # 五行 (text)
    wuxing = db.Column(db.Text, nullable=True, comment='五行')
    
    # 冲煞 (text)
    chongsha = db.Column(db.Text, nullable=True, comment='冲煞')
    
    # 彭祖百忌 (text)
    baiji = db.Column(db.Text, nullable=True, comment='彭祖百忌')
    
    # 吉神宜趋 (text)
    jishen = db.Column(db.Text, nullable=True, comment='吉神宜趋')
    
    # 宜 (text)
    yi = db.Column(db.Text, nullable=True, comment='宜')
    
    # 凶神宜忌 (text)
    xiongshen = db.Column(db.Text, nullable=True, comment='凶神宜忌')
    
    # 忌 (text)
    ji = db.Column(db.Text, nullable=True, comment='忌')
    
    # 颜色 (varchar(10))
    color = db.Column(db.String(10), nullable=True, comment='颜色')
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def __repr__(self):
        return f"<CalendarInfo(id={self.id}, yangli='{self.yangli}', yinli='{self.yinli}')>"
    
    def to_dict(self) -> dict:
        """
        将模型对象转换为字典
        
        Returns:
            dict: 包含所有字段的字典
        """
        return {
            'id': self.id,
            'yangli': self.yangli.isoformat() if self.yangli is not None else None,
            'yinli': self.yinli,
            'wuxing': self.wuxing,
            'chongsha': self.chongsha,
            'baiji': self.baiji,
            'jishen': self.jishen,
            'yi': self.yi,
            'xiongshen': self.xiongshen,
            'ji': self.ji,
            'color': self.color
        }
    
    def to_json(self) -> str:
        """
        将模型对象转换为JSON字符串
        
        Returns:
            str: JSON格式的字符串
        """
        return json.dumps(self.to_dict(), ensure_ascii=False, indent=2)
    
    @classmethod
    def from_dict(cls, data: dict) -> 'CalendarInfo':
        """
        从字典创建模型对象
        
        Args:
            data (dict): 包含字段数据的字典
            
        Returns:
            CalendarInfo: 新创建的模型对象
        """
        # 处理日期字段
        yangli = data.get('yangli')
        if isinstance(yangli, str):
            try:
                yangli = date.fromisoformat(yangli)
            except ValueError:
                yangli = None
        
        # 从wuxing字段中判断五行元素并设置对应的颜色值
        wuxing = data.get('wuxing', '') or ''
        color = data.get('color')  # 先获取字典中的color值
        
        # 如果字典中没有color值，则根据wuxing字段判断五行元素设置颜色
        if not color:
            if '金' in wuxing:
                color = '白'
            elif '水' in wuxing:
                color = '黑'
            elif '木' in wuxing:
                color = '绿'
            elif '火' in wuxing:
                color = '红'
            elif '土' in wuxing:
                color = '黄'
        
        return cls(
            yangli=yangli,  # type: ignore
            yinli=data.get('yinli'),  # type: ignore
            wuxing=wuxing,  # type: ignore
            chongsha=data.get('chongsha'),  # type: ignore
            baiji=data.get('baiji'),  # type: ignore
            jishen=data.get('jishen'),  # type: ignore
            yi=data.get('yi'),  # type: ignore
            xiongshen=data.get('xiongshen'),  # type: ignore
            ji=data.get('ji'),  # type: ignore
            color=color  # type: ignore
        )


class CalendarService:
    """
    黄历信息服务类
    提供黄历信息的增删改查操作
    """
    
    def __init__(self, engine=None):
        """
        初始化服务
        
        Args:
            engine: SQLAlchemy引擎对象，如果为None则使用Flask-SQLAlchemy的db.session
        """
        self.engine = engine
    
    def create_calendar_info(self, calendar_data: dict) -> CalendarInfo:
        """
        创建新的黄历信息记录
        
        Args:
            calendar_data (dict): 黄历信息数据
            
        Returns:
            CalendarInfo: 创建的黄历信息对象
        """
        calendar_info = CalendarInfo.from_dict(calendar_data)
        
        if self.engine:
            with Session(self.engine) as session:
                session.add(calendar_info)
                session.commit()
                session.refresh(calendar_info)
        else:
            # 使用Flask-SQLAlchemy的db.session
            db.session.add(calendar_info)
            db.session.commit()
            db.session.refresh(calendar_info)
            
        return calendar_info
    
    def get_calendar_by_date(self, yangli_date: date) -> Optional[CalendarInfo]:
        """
        根据阳历日期查询黄历信息
        
        Args:
            yangli_date (date): 阳历日期
            
        Returns:
            Optional[CalendarInfo]: 黄历信息对象，如果不存在则返回None
        """
        if self.engine:
            with Session(self.engine) as session:
                return session.query(CalendarInfo).filter(
                    CalendarInfo.yangli == yangli_date
                ).first()
        else:
            # 使用Flask-SQLAlchemy的db.session
            return CalendarInfo.query.filter(
                CalendarInfo.yangli == yangli_date
            ).first()
    
    def get_calendar_by_id(self, calendar_id: int) -> Optional[CalendarInfo]:
        """
        根据ID查询黄历信息
        
        Args:
            calendar_id (int): 黄历信息ID
            
        Returns:
            Optional[CalendarInfo]: 黄历信息对象，如果不存在则返回None
        """
        if self.engine:
            with Session(self.engine) as session:
                return session.query(CalendarInfo).filter(
                    CalendarInfo.id == calendar_id
                ).first()
        else:
            # 使用Flask-SQLAlchemy的db.session
            return CalendarInfo.query.filter(
                CalendarInfo.id == calendar_id
            ).first()
    
    def update_calendar_info(self, calendar_id: int, update_data: dict) -> Optional[CalendarInfo]:
        """
        更新黄历信息
        
        Args:
            calendar_id (int): 黄历信息ID
            update_data (dict): 要更新的数据
            
        Returns:
            Optional[CalendarInfo]: 更新后的黄历信息对象，如果不存在则返回None
        """
        if self.engine:
            with Session(self.engine) as session:
                calendar_info = session.query(CalendarInfo).filter(
                    CalendarInfo.id == calendar_id
                ).first()
                
                if not calendar_info:
                    return None
                
                # 更新字段
                for key, value in update_data.items():
                    if hasattr(calendar_info, key):
                        if key == 'yangli' and isinstance(value, str):
                            try:
                                value = date.fromisoformat(value)
                            except ValueError:
                                continue
                        setattr(calendar_info, key, value)
                
                session.commit()
                session.refresh(calendar_info)
                
        else:
            # 使用Flask-SQLAlchemy的db.session
            calendar_info = CalendarInfo.query.filter(
                CalendarInfo.id == calendar_id
            ).first()
            
            if not calendar_info:
                return None
            
            # 更新字段
            for key, value in update_data.items():
                if hasattr(calendar_info, key):
                    if key == 'yangli' and isinstance(value, str):
                        try:
                            value = date.fromisoformat(value)
                        except ValueError:
                            continue
                    setattr(calendar_info, key, value)
            
            db.session.commit()
            db.session.refresh(calendar_info)
            
        return calendar_info
    
    def delete_calendar_info(self, calendar_id: int) -> bool:
        """
        删除黄历信息
        
        Args:
            calendar_id (int): 黄历信息ID
            
        Returns:
            bool: 删除成功返回True，否则返回False
        """
        if self.engine:
            with Session(self.engine) as session:
                calendar_info = session.query(CalendarInfo).filter(
                    CalendarInfo.id == calendar_id
                ).first()
                
                if not calendar_info:
                    return False
                
                session.delete(calendar_info)
                session.commit()
        else:
            # 使用Flask-SQLAlchemy的db.session
            calendar_info = CalendarInfo.query.filter(
                CalendarInfo.id == calendar_id
            ).first()
            
            if not calendar_info:
                return False
            
            db.session.delete(calendar_info)
            db.session.commit()
            
        return True
    
    def get_calendar_list(self, limit: int = 100, offset: int = 0) -> list[CalendarInfo]:
        """
        获取黄历信息列表
        
        Args:
            limit (int): 限制返回数量，默认100
            offset (int): 偏移量，默认0
            
        Returns:
            list[CalendarInfo]: 黄历信息对象列表
        """
        if self.engine:
            with Session(self.engine) as session:
                return session.query(CalendarInfo).order_by(
                    CalendarInfo.yangli.desc()
                ).offset(offset).limit(limit).all()
        else:
            # 使用Flask-SQLAlchemy的db.session
            return CalendarInfo.query.order_by(
                CalendarInfo.yangli.desc()
            ).offset(offset).limit(limit).all()
    
    def search_calendar_by_keyword(self, keyword: str, limit: int = 100) -> list[CalendarInfo]:
        """
        根据关键词搜索黄历信息
        
        Args:
            keyword (str): 搜索关键词
            limit (int): 限制返回数量，默认100
            
        Returns:
            list[CalendarInfo]: 匹配的黄历信息对象列表
        """
        if self.engine:
            with Session(self.engine) as session:
                return session.query(CalendarInfo).filter(
                    (CalendarInfo.yinli.contains(keyword)) |
                    (CalendarInfo.wuxing.contains(keyword)) |
                    (CalendarInfo.chongsha.contains(keyword)) |
                    (CalendarInfo.baiji.contains(keyword)) |
                    (CalendarInfo.jishen.contains(keyword)) |
                    (CalendarInfo.yi.contains(keyword)) |
                    (CalendarInfo.xiongshen.contains(keyword)) |
                    (CalendarInfo.ji.contains(keyword)) |
                    (CalendarInfo.color.contains(keyword))
                ).limit(limit).all()
        else:
            # 使用Flask-SQLAlchemy的db.session
            return CalendarInfo.query.filter(
                (CalendarInfo.yinli.contains(keyword)) |
                (CalendarInfo.wuxing.contains(keyword)) |
                (CalendarInfo.chongsha.contains(keyword)) |
                (CalendarInfo.baiji.contains(keyword)) |
                (CalendarInfo.jishen.contains(keyword)) |
                (CalendarInfo.yi.contains(keyword)) |
                (CalendarInfo.xiongshen.contains(keyword)) |
                (CalendarInfo.ji.contains(keyword)) |
                (CalendarInfo.color.contains(keyword))
            ).limit(limit).all()
    
    def fetch_calendar_from_api(self, yangli_date: date) -> Optional[dict]:
        """
        从外部API获取黄历信息
        
        Args:
            yangli_date (date): 阳历日期
            
        Returns:
            Optional[dict]: API返回的黄历信息，如果失败则返回None
        """
        try:
            # 从配置文件获取API配置
            api_url = CALENDAR_API_CONFIG['url']
            api_key = CALENDAR_API_CONFIG['key']
            timeout = CALENDAR_API_CONFIG['timeout']
            
            # 格式化日期为YYYYMMDD格式
            date_str = yangli_date.strftime('%Y%m%d')
            
            # 请求参数
            request_params = {
                'key': api_key,
                'date': date_str,
            }
            
            # 发起API请求
            response = requests.get(api_url, params=request_params, timeout=timeout)
            
            if response.status_code == 200:
                response_result = response.json()
                
                # 检查API返回结果
                if response_result.get('error_code') == 0 and response_result.get('reason') == 'successed':
                    return response_result.get('result')
                else:
                    print(f"API返回错误: {response_result}")
                    return None
            else:
                print(f"API请求失败，状态码: {response.status_code}")
                return None
                
        except requests.exceptions.RequestException as e:
            print(f"API请求异常: {e}")
            return None
        except Exception as e:
            print(f"获取API数据时发生错误: {e}")
            return None
    
    def save_calendar_from_api(self, api_data: dict) -> Optional[CalendarInfo]:
        """
        将API返回的黄历信息保存到数据库
        
        Args:
            api_data (dict): API返回的黄历信息数据
            
        Returns:
            Optional[CalendarInfo]: 保存后的黄历信息对象，如果失败则返回None
        """
        try:
            # 解析API数据
            yangli_str = api_data.get('yangli')
            if not yangli_str:
                print("API数据中缺少阳历日期")
                return None
            
            # 解析日期
            try:
                yangli_date = date.fromisoformat(yangli_str)
            except ValueError:
                print(f"无效的日期格式: {yangli_str}")
                return None
            
            # 从wuxing字段中判断五行元素并设置对应的颜色值
            wuxing = api_data.get('wuxing', '') or ''
            color = api_data.get('color')  # 先获取API中的color值
            
            # 如果API中没有color值，则根据wuxing字段判断五行元素设置颜色
            if not color:
                if '金' in wuxing:
                    color = '白'
                elif '水' in wuxing:
                    color = '黑'
                elif '木' in wuxing:
                    color = '绿'
                elif '火' in wuxing:
                    color = '红'
                elif '土' in wuxing:
                    color = '黄'
            
            # 创建CalendarInfo对象
            calendar_info = CalendarInfo(
                yangli=yangli_date,  # type: ignore
                yinli=api_data.get('yinli', ''),  # type: ignore
                wuxing=wuxing,  # type: ignore
                chongsha=api_data.get('chongsha'),  # type: ignore
                baiji=api_data.get('baiji'),  # type: ignore
                jishen=api_data.get('jishen'),  # type: ignore
                yi=api_data.get('yi'),  # type: ignore
                xiongshen=api_data.get('xionshen'),  # type: ignore  # 注意API返回的是xionshen
                ji=api_data.get('ji'),  # type: ignore
                color=color  # type: ignore
            )
            
            # 保存到数据库
            if self.engine:
                with Session(self.engine) as session:
                    session.add(calendar_info)
                    session.commit()
                    session.refresh(calendar_info)
            else:
                # 使用Flask-SQLAlchemy的db.session
                db.session.add(calendar_info)
                db.session.commit()
                db.session.refresh(calendar_info)
                
            print(f"成功保存黄历信息到数据库，ID: {calendar_info.id}")
            return calendar_info
            
        except Exception as e:
            print(f"保存API数据到数据库时发生错误: {e}")
            return None


# 便捷函数
def create_calendar_info(calendar_data: dict, engine=None) -> CalendarInfo:
    """
    创建黄历信息的便捷函数
    
    Args:
        calendar_data (dict): 黄历信息数据
        engine: SQLAlchemy引擎对象
        
    Returns:
        CalendarInfo: 创建的黄历信息对象
    """
    service = CalendarService(engine)
    return service.create_calendar_info(calendar_data)


def get_calendar_by_date(yangli_date: str, engine=None) -> dict:
    """
    根据阳历日期查询黄历信息的便捷函数
    
    Args:
        yangli_date (str): 阳历日期，格式为YYYY-MM-DD
        engine: SQLAlchemy引擎对象
        
    Returns:
        dict: 包含查询结果的JSON格式数据
    """
    try:
        # 验证日期格式
        if not isinstance(yangli_date, str) or len(yangli_date) != 10:
            return {
                "reason": "failed",
                "return_code": 400,
                "result": None,
                "error": "日期格式错误，请使用YYYY-MM-DD格式"
            }
        
        # 解析日期字符串
        try:
            parsed_date = date.fromisoformat(yangli_date)
        except ValueError:
            return {
                "reason": "failed",
                "return_code": 400,
                "result": None,
                "error": "无效的日期格式"
            }
        
        # 查询数据库
        service = CalendarService(engine)
        calendar_info = service.get_calendar_by_date(parsed_date)
        
        if calendar_info:
            # 查询成功，返回指定格式的JSON数据
            return {
                "reason": "successed",
                "return_code": 200,
                "result": {
                    "id": str(calendar_info.id),
                    "yangli": calendar_info.yangli.isoformat() if calendar_info.yangli is not None else None,
                    "yinli": calendar_info.yinli,
                    "wuxing": calendar_info.wuxing,
                    "chongsha": calendar_info.chongsha,
                    "baiji": calendar_info.baiji,
                    "jishen": calendar_info.jishen,
                    "yi": calendar_info.yi,
                    "xiongshen": calendar_info.xiongshen,
                    "ji": calendar_info.ji,
                    "color": calendar_info.color
                }
            }
        else:
            # 数据库中没有找到记录，尝试从API获取
            print(f"数据库中没有找到日期 {yangli_date} 的黄历信息，尝试从API获取...")
            
            # 从API获取数据
            api_data = service.fetch_calendar_from_api(parsed_date)
            
            if api_data:
                # API获取成功，保存到数据库
                print("API获取数据成功，正在保存到数据库...")
                saved_calendar = service.save_calendar_from_api(api_data)
                
                if saved_calendar:
                    # 保存成功，返回数据
                    return {
                        "reason": "successed",
                        "return_code": 200,
                        "result": {
                            "id": str(saved_calendar.id),
                            "yangli": saved_calendar.yangli.isoformat() if saved_calendar.yangli is not None else None,
                            "yinli": saved_calendar.yinli,
                            "wuxing": saved_calendar.wuxing,
                            "chongsha": saved_calendar.chongsha,
                            "baiji": saved_calendar.baiji,
                            "jishen": saved_calendar.jishen,
                            "yi": saved_calendar.yi,
                            "xiongshen": saved_calendar.xiongshen,
                            "ji": saved_calendar.ji,
                            "color": saved_calendar.color
                        }
                    }
                else:
                    # 保存到数据库失败
                    return {
                        "reason": "failed",
                        "return_code": 500,
                        "result": None,
                        "error": f"API获取数据成功但保存到数据库失败"
                    }
            else:
                # API获取失败
                return {
                    "reason": "failed",
                    "return_code": 404,
                    "result": None,
                    "error": f"未找到日期 {yangli_date} 的黄历信息，且API获取失败"
                }
            
    except Exception as e:
        # 发生异常
        return {
            "reason": "failed",
            "return_code": 500,
            "result": None,
            "error": f"查询过程中发生错误: {str(e)}"
        }


def get_calendar_by_id(calendar_id: int, engine=None) -> Optional[CalendarInfo]:
    """
    根据ID查询黄历信息的便捷函数
    
    Args:
        calendar_id (int): 黄历信息ID
        engine: SQLAlchemy引擎对象
        
    Returns:
        Optional[CalendarInfo]: 黄历信息对象
    """
    service = CalendarService(engine)
    return service.get_calendar_by_id(calendar_id)


# 导出主要类和函数
__all__ = [
    'CalendarInfo',
    'CalendarService',
    'create_calendar_info',
    'get_calendar_by_date',
    'get_calendar_by_id'
]