import logging
import os
import sys
import json
from typing import Dict, List, Optional, Any, Union
from pydantic import BaseModel, Field
from google import genai
from google.genai.types import Tool, GenerateContentConfig, GoogleSearch

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

# 定义基础数据结构
class TextAnalysisResult(BaseModel):
    """文本分析结果，包含分析正文、推理过程和引用URL"""
    analysis_text: str = Field(description="详细的分析文本")
    reasoning_process: Optional[str] = Field(description="模型的推理过程", default=None)
    references: Optional[List[str]] = Field(description="参考资料和引用URL列表", default=None)

class NumericalAnalysisResult(BaseModel):
    """数值分析结果，包含数值和分析描述"""
    value: str = Field(description="评估值")
    description: str = Field(description="评估描述")

# 分析类型到模型类的映射
ANALYSIS_TYPE_MODEL_MAP = {
    # 文本分析类型
    "text_analysis_result": TextAnalysisResult,
    
    # 数值分析类型
    "numerical_analysis_result": NumericalAnalysisResult
}

class JsonOutputParser:
    """解析JSON输出的工具类"""
    def __init__(self, pydantic_object):
        """初始化解析器
        
        Args:
            pydantic_object: 需要解析成的Pydantic模型类
        """
        self.pydantic_object = pydantic_object

    def get_format_instructions(self) -> str:
        """获取格式化指令"""
        schema = self.pydantic_object.schema()
        schema_str = json.dumps(schema, ensure_ascii=False, indent=2)
        
        return f"""请将你的回答格式化为符合以下JSON结构的格式:
{schema_str}

你的回答应该只包含一个有效的JSON对象，而不包含其他内容。请不要包含任何解释或前缀，仅输出JSON本身。
"""

    def parse(self, text: str) -> Any:
        """解析文本为对象"""
        try:
            # 尝试解析JSON
            text = text.strip()
            # 如果文本以```json开头且以```结尾，则提取中间部分
            if text.startswith("```json") and text.endswith("```"):
                text = text[7:-3].strip()
            # 如果文本以```开头且以```结尾，则提取中间部分
            elif text.startswith("```") and text.endswith("```"):
                text = text[3:-3].strip()
                
            json_obj = json.loads(text)
            return self.pydantic_object.parse_obj(json_obj)
        except Exception as e:
            raise ValueError(f"无法解析为JSON或无法匹配模式: {e}")

# 获取项目根目录的绝对路径
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# 尝试导入配置
try:
    # 添加项目根目录到 Python 路径
    sys.path.append(os.path.dirname(ROOT_DIR))
    sys.path.append(ROOT_DIR)

    # 导入配置
    try:
        from scripts.config import get_model_config

        logger.info("成功从scripts.config导入配置")
    except ImportError:
        try:
            from src.scripts.config import get_model_config

            logger.info("成功从src.scripts.config导入配置")
        except ImportError:
            logger.warning("无法导入配置模块，使用默认配置")


            # 使用默认配置的实现
            def get_model_config(platform: str, model_type: str):
                return {
                    "api_key": "AIzaSyAVE8yTaPtN-TxCCHTc9Jb-aCV-Xo1EFuU",
                    "model": "gemini-2.0-flash"
                }
except Exception as e:
    logger.error(f"导入配置时出错: {str(e)}，使用默认配置")


    # 使用默认配置的实现
    def get_model_config(platform: str, model_type: str):
        return {
            "api_key": "AIzaSyAVE8yTaPtN-TxCCHTc9Jb-aCV-Xo1EFuU",
            "model": "gemini-2.0-flash"
        }

class ChatBot:
    def __init__(self, platform: str = "Gemini", model_type: str = "gemini-2.0-flash"):
        """初始化聊天机器人

        Args:
            platform: 平台名称，默认为Gemini
            model_type: 要使用的模型类型，默认为gemini-2.0-flash
        """
        try:
            # 从配置获取API配置
            config = get_model_config(platform, model_type)
            self.api_key = config["api_key"]
            self.model = config["model"]  # 直接使用配置中的模型名称

            logger.info(f"初始化ChatBot，使用平台: {platform}, 模型: {self.model}")

            google_search_tool = Tool(
                google_search=GoogleSearch()
            )

            self.client = genai.Client(api_key=self.api_key)
            
            # 系统提示语
            self.system_instruction = """你是一位经验丰富的专业投资经理，擅长基本面分析和投资决策。你的分析特点如下：

                            1. 分析风格：
                            - 专业、客观、理性
                            - 注重数据支撑
                            - 关注风险控制
                            - 重视投资性价比

                            2. 分析框架：
                            - 公司基本面分析
                            - 行业竞争格局
                            - 估值水平评估
                            - 风险因素识别
                            - 投资机会判断

                            3. 输出要求：
                            - 简明扼要，重点突出
                            - 逻辑清晰，层次分明
                            - 数据准确，论据充分
                            - 结论明确，建议具体

                            请用专业投资经理的视角，对股票进行深入分析和投资建议。如果信息不足，请明确指出。"""
            
            # 对话历史
            self.conversation_history = []
        except Exception as e:
            logger.error(f"初始化ChatBot时出错: {str(e)}")
            raise

    def chat(self, user_input: str, analysis_type: Optional[str] = None, temperature: float = 0.7, top_p: float = 0.7, max_tokens: int = 4096) -> dict:
        """处理用户输入并返回AI回复
        
        Args:
            user_input: 用户输入的问题
            analysis_type: 分析类型，如company_profile或business_scale等
            temperature: 控制输出的随机性，范围0-2，默认0.7
            top_p: 控制输出的多样性，范围0-1，默认0.7
            max_tokens: 控制输出的最大长度，默认4096
            
        Returns:
            dict: 包含AI回复的字典
        """
        try:
            from google.genai import types
            google_search_tool = Tool(
                google_search=GoogleSearch()
            )
            
            # 确定分析类型和模型
            model_class = None
            if analysis_type and analysis_type in ANALYSIS_TYPE_MODEL_MAP:
                model_class = ANALYSIS_TYPE_MODEL_MAP[analysis_type]
                parser = JsonOutputParser(pydantic_object=model_class)
                format_instructions = parser.get_format_instructions()
                
                # 在原有基础上添加格式指令
                combined_instruction = f"{self.system_instruction}\n\n{format_instructions}"
            else:
                combined_instruction = self.system_instruction
                parser = None
            
            response = self.client.models.generate_content_stream(
                model=self.model,
                config=GenerateContentConfig(
                    system_instruction=combined_instruction,
                    tools=[google_search_tool],
                    response_modalities=["TEXT"],
                    temperature=temperature,
                    top_p=top_p,
                    max_output_tokens=max_tokens
                ),
                contents=user_input
            )
            
            full_response = ""
            for chunk in response:
                full_response += chunk.text
                print(chunk.text, end="")
            
            # 如果指定了分析类型，尝试解析输出
            if parser:
                try:
                    # 检查分析类型是文本分析还是数值分析
                    is_text_analysis = model_class == TextAnalysisResult
                    is_numerical_analysis = model_class == NumericalAnalysisResult
                    
                    # 使用解析方法尝试处理JSON响应
                    json_result = self.parse_json_response(full_response, model_class)
                    if json_result:
                        return json_result
                    
                    # 如果使用解析方法失败，尝试使用parser解析
                    parsed_result = parser.parse(full_response)
                    structured_result = parsed_result.dict()
                    
                    if isinstance(parsed_result, TextAnalysisResult):
                        # 文本分析结果
                        return {
                            "response": structured_result["analysis_text"],
                            "reasoning_process": structured_result.get("reasoning_process"),
                            "references": structured_result.get("references")
                        }
                    elif isinstance(parsed_result, NumericalAnalysisResult):
                        # 数值分析结果
                        return {
                            "response": structured_result["value"],
                            "description": structured_result.get("description", "")
                        }
                except Exception as e:
                    logger.warning(f"解析响应失败: {str(e)}，返回原始响应")
                    
                    # 最后尝试根据分析类型直接构建结果
                    if is_text_analysis:
                        return {
                            "response": full_response,
                            "reasoning_process": None,
                            "references": None
                        }
                    elif is_numerical_analysis:
                        try:
                            # 尝试从文本中提取数字
                            import re
                            number_match = re.search(r'\b(\d+)\b', full_response)
                            value = int(number_match.group(1)) if number_match else 0
                            return {
                                "response": value,
                                "description": full_response
                            }
                        except:
                            return {
                                "response": 0,
                                "description": full_response
                            }
                    else:
                        return {
                            "response": full_response,
                            "reasoning_process": None,
                            "references": None
                        }
            else:
                # 普通响应
                return {
                    "response": full_response,
                    "reasoning_process": None,
                    "references": None
                }

        except Exception as e:
            logger.error(f"聊天失败: {str(e)}")
            error_msg = f"抱歉，发生错误: {str(e)}"
            print(f"\n{error_msg}")
            return {"response": error_msg, "reasoning_process": None, "references": None}

    def clear_history(self):
        """清除对话历史"""
        self.conversation_history = []
        # 重置会话
        if hasattr(self, 'model_instance'):
            self.chat_session = self.model_instance.start_chat()
        print("对话历史已清除")

    def run(self):
        """运行聊天机器人"""
        print("欢迎使用Gemini AI助手！输入 'quit' 退出，输入 'clear' 清除对话历史。")
        print("-" * 50)

        while True:
            try:
                # 获取用户输入
                user_input = input("\n你: ").strip()

                # 检查退出命令
                if user_input.lower() == 'quit':
                    print("感谢使用，再见！")
                    break

                # 检查清除历史命令
                if user_input.lower() == 'clear':
                    self.clear_history()
                    continue

                # 如果输入为空，继续下一轮
                if not user_input:
                    continue

                # 获取AI回复
                self.chat(user_input)
                print("-" * 50)

            except KeyboardInterrupt:
                print("\n感谢使用，再见！")
                break
            except Exception as e:
                logger.error(f"运行错误: {str(e)}")
                print(f"发生错误: {str(e)}")

    def parse_json_response(self, full_response: str, model_class) -> dict:
        """解析JSON响应，处理不同的JSON结构
        
        Args:
            full_response: 完整的响应文本
            model_class: 模型类（TextAnalysisResult或NumericalAnalysisResult）
            
        Returns:
            dict: 包含解析结果的字典
        """
        # 检查分析类型
        is_text_analysis = model_class == TextAnalysisResult
        is_numerical_analysis = model_class == NumericalAnalysisResult
        
        try:
            # 尝试提取JSON内容
            json_text = full_response
            # 如果文本以```json开头且以```结尾，则提取中间部分
            if "```json" in json_text and "```" in json_text.split("```json", 1)[1]:
                json_text = json_text.split("```json", 1)[1].split("```", 1)[0].strip()
            # 如果文本以```开头且以```结尾，则提取中间部分
            elif "```" in json_text and "```" in json_text.split("```", 1)[1]:
                json_text = json_text.split("```", 1)[1].split("```", 1)[0].strip()
            
            # 尝试解析JSON
            json_obj = json.loads(json_text)
            
            # 情况1: 直接包含字段的格式 {"analysis_text": "...", "reasoning_process": null, "references": [...]}
            if is_text_analysis and "analysis_text" in json_obj:
                return {
                    "response": json_obj["analysis_text"],
                    "reasoning_process": json_obj.get("reasoning_process"),
                    "references": json_obj.get("references", [])
                }
            elif is_numerical_analysis and "value" in json_obj:
                value = json_obj["value"]
                if isinstance(value, str) and value.isdigit():
                    value = int(value)
                elif isinstance(value, (int, float)):
                    value = int(value)
                else:
                    value = 0
                    
                return {
                    "response": value,
                    "description": json_obj.get("description", "")
                }
            
            # 情况2: properties中包含字段的格式
            if "properties" in json_obj and isinstance(json_obj["properties"], dict):
                properties = json_obj["properties"]
                
                if is_text_analysis and "analysis_text" in properties:
                    return {
                        "response": properties["analysis_text"],
                        "reasoning_process": properties.get("reasoning_process"),
                        "references": properties.get("references", [])
                    }
                elif is_numerical_analysis and "value" in properties:
                    value = properties["value"]
                    if isinstance(value, str) and value.isdigit():
                        value = int(value)
                    elif isinstance(value, (int, float)):
                        value = int(value)
                    else:
                        value = 0
                        
                    return {
                        "response": value,
                        "description": properties.get("description", "")
                    }
            
            # 如果无法识别结构，使用parser解析
            raise ValueError("无法识别的JSON结构，尝试使用parser解析")
            
        except Exception as e:
            # 记录错误但不抛出，让调用方法继续尝试其他解析方式
            logger.warning(f"JSON解析失败: {str(e)}")
            return None

if __name__ == "__main__":
    # 设置日志级别
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    # 创建并运行聊天机器人
    bot = ChatBot()
    bot.run() 