stock-monitor/app/dao/watchlist_dao.py

"""
监控列表数据访问对象
"""
from typing import Dict, List, Optional, Tuple
from datetime import datetime

from .base_dao import BaseDAO


class WatchlistDAO(BaseDAO):
    """监控列表数据访问对象"""

    def get_watchlist(self) -> List[Dict]:
        """获取完整的监控列表，包含股票信息"""
        query = """
        SELECT
            w.stock_code,
            s.stock_name,
            s.market,
            w.target_market_value_min,
            w.target_market_value_max,
            w.created_at,
            w.updated_at
        FROM watchlist w
        JOIN stocks s ON w.stock_code = s.stock_code
        ORDER BY w.created_at DESC
        """
        return self._execute_query(query)

    def add_to_watchlist(self, stock_code: str, target_min: float = None,
                        target_max: float = None) -> bool:
        """添加股票到监控列表"""
        try:
            # 检查是否已在监控列表中
            existing = self.get_watchlist_item(stock_code)
            if existing:
                # 更新现有的目标市值
                return self.update_watchlist_item(stock_code, target_min, target_max)

            # 添加新项到监控列表
            query = """
            INSERT INTO watchlist (stock_code, target_market_value_min, target_market_value_max)
            VALUES (%s, %s, %s)
            """
            self._execute_insert(query, (stock_code, target_min, target_max))

            self.log_data_update('watchlist', stock_code, 'success', 'Added to watchlist')
            return True

        except Exception as e:
            self.logger.error(f"添加到监控列表失败: {stock_code}, 错误: {e}")
            self.log_data_update('watchlist', stock_code, 'failed', str(e))
            return False

    def remove_from_watchlist(self, stock_code: str) -> bool:
        """从监控列表移除股票"""
        try:
            query = "DELETE FROM watchlist WHERE stock_code = %s"
            affected_rows = self._execute_update(query, (stock_code,))
            success = affected_rows > 0

            if success:
                self.log_data_update('watchlist', stock_code, 'success', 'Removed from watchlist')
            else:
                self.log_data_update('watchlist', stock_code, 'failed', 'Stock not found in watchlist')

            return success

        except Exception as e:
            self.logger.error(f"从监控列表移除失败: {stock_code}, 错误: {e}")
            self.log_data_update('watchlist', stock_code, 'failed', str(e))
            return False

    def get_watchlist_item(self, stock_code: str) -> Optional[Dict]:
        """获取监控列表中的单个项目"""
        query = """
        SELECT
            w.stock_code,
            s.stock_name,
            s.market,
            w.target_market_value_min,
            w.target_market_value_max,
            w.created_at,
            w.updated_at
        FROM watchlist w
        JOIN stocks s ON w.stock_code = s.stock_code
        WHERE w.stock_code = %s
        """
        return self._execute_single_query(query, (stock_code,))

    def update_watchlist_item(self, stock_code: str, target_min: float = None,
                             target_max: float = None) -> bool:
        """更新监控列表项目"""
        try:
            query = """
            UPDATE watchlist
            SET target_market_value_min = %s,
                target_market_value_max = %s,
                updated_at = CURRENT_TIMESTAMP
            WHERE stock_code = %s
            """
            affected_rows = self._execute_update(query, (target_min, target_max, stock_code))
            success = affected_rows > 0

            if success:
                self.log_data_update('watchlist', stock_code, 'success', 'Updated watchlist item')
            else:
                self.log_data_update('watchlist', stock_code, 'failed', 'Stock not found in watchlist')

            return success

        except Exception as e:
            self.logger.error(f"更新监控列表失败: {stock_code}, 错误: {e}")
            self.log_data_update('watchlist', stock_code, 'failed', str(e))
            return False

    def get_watchlist_with_data(self, data_date: str = None) -> List[Dict]:
        """获取监控列表及其股票数据"""
        if data_date is None:
            data_date = self.get_today_date()

        query = """
        SELECT
            w.stock_code,
            s.stock_name,
            s.market,
            w.target_market_value_min,
            w.target_market_value_max,
            sd.price,
            sd.change_percent,
            sd.market_value as current_market_value,
            sd.pe_ratio,
            sd.pb_ratio,
            sd.from_cache
        FROM watchlist w
        JOIN stocks s ON w.stock_code = s.stock_code
        LEFT JOIN stock_data sd ON w.stock_code = sd.stock_code AND sd.data_date = %s
        ORDER BY w.created_at DESC
        """
        return self._execute_query(query, (data_date,))

    def clear_watchlist(self) -> bool:
        """清空监控列表"""
        try:
            query = "DELETE FROM watchlist"
            self._execute_update(query)
            self.log_data_update('watchlist', 'all', 'success', 'Cleared watchlist')
            return True
        except Exception as e:
            self.logger.error(f"清空监控列表失败: {e}")
            self.log_data_update('watchlist', 'all', 'failed', str(e))
            return False

    def get_watchlist_count(self) -> int:
        """获取监控列表股票数量"""
        query = "SELECT COUNT(*) as count FROM watchlist"
        result = self._execute_single_query(query)
        return result['count'] if result else 0

    def get_stocks_needing_update(self, data_date: str = None) -> List[str]:
        """获取需要更新数据的股票代码列表"""
        if data_date is None:
            data_date = self.get_today_date()

        query = """
        SELECT DISTINCT w.stock_code
        FROM watchlist w
        LEFT JOIN stock_data sd ON w.stock_code = sd.stock_code AND sd.data_date = %s
        WHERE sd.stock_code IS NULL OR sd.data_date < %s
        """
        results = self._execute_query(query, (data_date, data_date))
        return [item['stock_code'] for item in results]