import logging
from typing import List, Dict, Tuple
from fundamentals_llm.fundamental_analysis_database import get_db, get_analysis_result
from sqlalchemy import text

# 设置日志记录
logger = logging.getLogger(__name__)

class EnterpriseScreener:
    """企业筛选器"""
    
    def __init__(self):
        """初始化筛选器"""
        self.db = next(get_db())
    
    def _get_stock_info(self, stock_code: str) -> Tuple[str, str]:
        """获取股票信息
        
        Args:
            stock_code: 股票代码
            
        Returns:
            Tuple[str, str]: (股票代码, 股票名称)
        """
        try:
            # 从任意维度获取股票信息
            result = get_analysis_result(self.db, stock_code, "company_profile")
            if result:
                return stock_code, result.stock_name
            return stock_code, "未知"
        except Exception as e:
            logger.error(f"获取股票信息失败: {str(e)}")
            return stock_code, "未知"

    def _screen_stocks_by_conditions(self, conditions: List[Dict], limit: int = None) -> List[Tuple[str, str]]:
        """通用的股票筛选方法
        
        Args:
            conditions: 筛选条件列表，每个条件是一个字典，包含：
                - dimension: 分析维度
                - field: extra_info中的字段名
                - operator: 操作符 ('=', '!=', '>=', '<=', 'in')
                - value: 比较值
                - required: 是否必须满足（默认True）
            limit: 限制返回结果数量
        
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        try:
            # 构建SQL查询条件
            where_conditions = []
            having_conditions = []
            for i, condition in enumerate(conditions):
                dimension = condition['dimension']
                field = condition['field']
                operator = condition['operator']
                value = condition['value']
                required = condition.get('required', True)
                
                # 构建JSON提取表达式
                json_extract = f"JSON_EXTRACT(fa.extra_info, '$.{field}')"
                if operator in ('=', '!='):
                    json_extract = f"JSON_UNQUOTE({json_extract})"
                elif operator in ('>=', '<=', '>', '<'):
                    json_extract = f"CAST({json_extract} AS SIGNED)"
                
                # 构建比较条件
                if operator == 'in':
                    value_str = ', '.join(str(v) for v in value)
                    where_condition = f"{json_extract} IN ({value_str})"
                else:
                    where_condition = f"{json_extract} {operator} '{value}'"
                
                # 添加到WHERE子句
                where_conditions.append(f"(fa.dimension = '{dimension}' AND {where_condition})")
                
                # 添加到HAVING子句
                if required:
                    having_conditions.append(f"WHEN fa.dimension = '{dimension}' AND {where_condition} THEN {i+1}")
            
            # 构建完整的SQL查询
            sql = f"""
            SELECT fa.stock_code, fa.stock_name
            FROM fundamental_analysis fa
            WHERE 1=1
                AND ({' OR '.join(where_conditions)})
            GROUP BY fa.stock_code, fa.stock_name
            HAVING COUNT(DISTINCT CASE
                {' '.join(having_conditions)}
            END) = {len(having_conditions)}
            """
            
            # 添加LIMIT子句
            if limit is not None:
                sql += f"\nLIMIT {limit}"
            
            # 打印SQL查询语句到控制台
            print("\n===================== SQL查询 =====================")
            print(sql)
            print("==================================================\n")
            
            # 执行查询
            cursor = self.db.execute(text(sql))
            qualified_stocks = []
            for row in cursor:
                stock_code = row[0]
                stock_name = row[1] if row[1] != None else self._get_stock_info(stock_code)[1]
                qualified_stocks.append((stock_code, stock_name))
            
            print(f"查询结果: 找到 {len(qualified_stocks)} 个符合条件的股票")
            
            return qualified_stocks
            
        except Exception as e:
            logger.error(f"通用筛选方法失败: {str(e)}")
            print(f"SQL查询执行失败: {str(e)}")
            return []

    def screen_multi_term_investment_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选推荐投资企业
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'investment_advice',
                'field': 'investment_advice_type',
                'operator': 'in',
                'value': ["'短期'","'中期'","'长期'"]
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': -1
            },
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '>=',
                'value': 0
            },
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_high_growth_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选高成长潜力企业
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '=',
                'value': 2
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'recent_projects',
                'field': 'major_events',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'odds',
                'operator': '=',
                'value': 1
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_stable_leaders(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选稳定型龙头企业
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'management_ownership',
                'field': 'ability',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'management_ownership',
                'field': 'shareholding',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 0
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_short_term_opportunities(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选短期投资机会
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'recent_projects',
                'field': 'major_events',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'stock_discussion',
                'field': 'emotion',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'odds',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 0
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_value_investment_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选价值型投资标的
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '=',
                'value': 2
            },
            {
                'dimension': 'management_ownership',
                'field': 'shareholding',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 0
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_turnaround_opportunities(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选困境反转机会
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '=',
                'value': -1
            },
            {
                'dimension': 'recent_projects',
                'field': 'major_events',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'stock_discussion',
                'field': 'emotion',
                'operator': '=',
                'value': 0
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '=',
                'value': 0
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_risk_averse_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选风险规避标的
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'management_ownership',
                'field': 'ability',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'management_ownership',
                'field': 'shareholding',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 0
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_innovation_driven_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选创新驱动型企业
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '=',
                'value': 2
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '=',
                'value': 2
            },
            {
                'dimension': 'recent_projects',
                'field': 'major_events',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 0
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 1
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_industry_integration_stocks(self, limit: int = None) -> List[Tuple[str, str]]:
        """筛选行业整合机会
        
        Args:
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        conditions = [
            {
                'dimension': 'industry_competition',
                'field': 'industry_space',
                'operator': '=',
                'value': 1
            },
            {
                'dimension': 'industry_cooperation',
                'field': 'collaboration_dynamics',
                'operator': '=',
                'value': 2
            },
            {
                'dimension': 'financial_report',
                'field': 'financial_report_level',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'securities_rating',
                'operator': '>=',
                'value': 1
            },
            {
                'dimension': 'target_price',
                'field': 'odds',
                'operator': '=',
                'value': 1
            }
        ]
        return self._screen_stocks_by_conditions(conditions, limit)

    def screen_by_custom_conditions(self, conditions: List[Dict], limit: int = None) -> List[Tuple[str, str]]:
        """根据自定义条件筛选企业
        
        Args:
            conditions: 筛选条件列表，每个条件是一个字典，包含：
                - dimension: 分析维度
                - field: extra_info中的字段名
                - operator: 操作符 ('=', '!=', '>=', '<=', 'in')
                - value: 比较值
                - required: 是否必须满足（默认True）
            limit: 限制返回结果数量
            
        Returns:
            List[Tuple[str, str]]: 符合条件的股票列表，每个元素为(股票代码, 股票名称)
        """
        print("\n================== 自定义条件筛选 ==================")
        print(f"条件数量: {len(conditions)}")
        for i, condition in enumerate(conditions):
            print(f"条件 {i+1}: 维度={condition.get('dimension')}, 字段={condition.get('field')}, " + 
                  f"操作符={condition.get('operator')}, 值={condition.get('value')}, " + 
                  f"必须满足={condition.get('required', True)}")
        print("====================================================\n")
        
        return self._screen_stocks_by_conditions(conditions, limit)

def main():
    """主函数"""
    # 设置日志级别
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建筛选器实例
    screener = EnterpriseScreener()
    
    # 测试选项
    print("\n请选择要筛选的企业类型：")
    print("1. 高成长潜力企业")
    print("2. 稳定型龙头企业")
    print("3. 短期投资机会")
    print("4. 价值型投资标的")
    print("5. 困境反转机会")
    print("6. 风险规避标的")
    print("7. 创新驱动型企业")
    print("8. 行业整合机会")
    print("9. 推荐投资企业")
    print("10. 退出")
    
    choice = input("\n请输入选项（1-10）: ").strip()
    
    if choice == "1":
        stocks = screener.screen_high_growth_stocks()
        print("\n高成长潜力企业：")
    elif choice == "2":
        stocks = screener.screen_stable_leaders()
        print("\n稳定型龙头企业：")
    elif choice == "3":
        stocks = screener.screen_short_term_opportunities()
        print("\n短期投资机会：")
    elif choice == "4":
        stocks = screener.screen_value_investment_stocks()
        print("\n价值型投资标的：")
    elif choice == "5":
        stocks = screener.screen_turnaround_opportunities()
        print("\n困境反转机会：")
    elif choice == "6":
        stocks = screener.screen_risk_averse_stocks()
        print("\n风险规避标的：")
    elif choice == "7":
        stocks = screener.screen_innovation_driven_stocks()
        print("\n创新驱动型企业：")
    elif choice == "8":
        stocks = screener.screen_industry_integration_stocks()
        print("\n行业整合机会：")
    elif choice == "9":
        stocks = screener.screen_multi_term_investment_stocks()
        print("\n推荐投资企业：")
    elif choice == "10":
        print("程序退出")
        return
    else:
        print("无效的选项")
        return
    
    # 输出结果
    if stocks:
        print(f"\n共找到 {len(stocks)} 只符合条件的股票：")
        for code, name in stocks:
            print(f"{name}({code})")
    else:
        print("\n未找到符合条件的股票")

if __name__ == "__main__":
    main() 