import os
from langchain_classic.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from adata import stock
from langchain_core.tools import tool
from telegram import Update
from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from telegram.error import TimedOut

load_dotenv()

# VectorDB 存在 ./rules_data 目录中
# ======================
# Embedding模型
# ======================

embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-small-zh-v1.5",
    model_kwargs={
        "device":"cpu"
    }
)

# ======================
# 加载向量库
# ======================

vector_store = Chroma(
    persist_directory="./rag_db",
    collection_name="investment_knowledge",
    embedding_function=embeddings
)

# 公司事实 Retriever
company_retriever = vector_store.as_retriever(
    search_kwargs={
        "k":5,
        "filter":{
            "category":"company_case"
        }
    }
)

# 行业 Retriever
industry_retriever = vector_store.as_retriever(
    search_kwargs={
        "k":3,
        "filter":{
            "category":"industry"
        }
    }
)

# 规则 Retriever
rule_retriever = vector_store.as_retriever(
    search_kwargs={
        "k":5,
        "filter":{
            "category":"rule"
        }
    }
)

def research_retrieve(question):

    result = {}
    # 公司事实
    company_docs = company_retriever.invoke(question)
    result["company"] = "\n\n".join(
        [
            doc.page_content
            for doc in company_docs
        ]
    )
    # 行业知识
    industry_docs = industry_retriever.invoke(question)
    result["industry"] = "\n\n".join(
        [
            doc.page_content
            for doc in industry_docs
        ]
    )
    # 分析规则
    rule_docs = rule_retriever.invoke(question)
    result["rules"] = "\n\n".join(
        [
            doc.page_content
            for doc in rule_docs
        ]
    )
    return result

@tool
def get_fina_info(stock_code: str) -> str:
    """获取指定股票代码的财务数据。当需要获取财务数据的时使用此工具"""

    # 字典定义：英文指标名 -> 中文财务指标名
    column_mapping = {
        # 基础与披露信息
        'stock_code': '股票代码',
        'short_name': '股票简称',
        'report_date': '报告期',
        'report_type': '报表类型',
        'notice_date': '公告日期',

        # 每股指标
        'basic_eps': '基本每股收益',
        'diluted_eps': '稀释每股收益',
        'non_gaap_eps': '扣非每股收益',
        'net_asset_ps': '每股净资产',
        'cap_reserve_ps': '每股公积金',
        'undist_profit_ps': '每股未分配利润',
        'oper_cf_ps': '每股经营现金流',

        # 规模与利润指标 (元)
        'total_rev': '营业总收入',
        'gross_profit': '毛利润',
        'net_profit_attr_sh': '归母净利润',
        'non_gaap_net_profit': '扣非归母净利润',

        # 增长率指标 (同比 YoY / 环比 QoQ)
        'total_rev_yoy_gr': '营业总收入同比增长率',
        'net_profit_yoy_gr': '归母净利润同比增长率',
        'non_gaap_net_profit_yoy_gr': '扣非归母净利润同比增长率',
        'total_rev_qoq_gr': '营业总收入环比增长率',
        'net_profit_qoq_gr': '归母净利润环比增长率',
        'non_gaap_net_profit_qoq_gr': '扣非归母净利润环比增长率',

        # 盈利能力与收益率
        'roe_wtd': '加权净资产收益率(ROE)',
        'roe_non_gaap_wtd': '扣非加权净资产收益率(ROE)',
        'roa_wtd': '加权总资产收益率(ROA)',
        'gross_margin': '销售毛利率',
        'net_margin': '销售净利率',

        # 现金流与收入质量比率
        'adv_receipts_to_rev': '预收款及合同负债占营收比',
        'net_cf_sales_to_rev': '销售收现率(销售现金流/营收)',
        'oper_cf_to_rev': '经营现金净流量占营收比',
        'eff_tax_rate': '实际有效税率',

        # 偿债能力与财务杠杆
        'curr_ratio': '流动比率',
        'quick_ratio': '速动比率',
        'cash_flow_ratio': '现金流量比率',
        'asset_liab_ratio': '资产负债率',
        'equity_multiplier': '权益乘数',
        'equity_ratio': '产权比率',

        # 运营效率与周转率/天数
        'total_asset_turn_days': '总资产周转天数',
        'inv_turn_days': '存货周转天数',
        'acct_recv_turn_days': '应收账款周转天数',
        'total_asset_turn_rate': '总资产周转率',
        'inv_turn_rate': '存货周转率',
        'acct_recv_turn_rate': '应收账款周转率'
    }

    df = stock.finance.get_core_index(stock_code=stock_code)

    year_data = df[df.report_type == "年报"].iloc[:5, :]  # 取5年的年报数据
    year_data.rename(columns=column_mapping, inplace=True)

    year_formatted_lines = []
    for _, row in year_data.iterrows():
        line_items = []
        for col, val in zip(year_data.columns, row):
            # 判断是否为数字类型
            if isinstance(val, (int, float)):
                # 如果列名包含"率"，则格式化为保留两位小数并加上 %
                if "率" in col:
                    # 如果原始数值本身就是小数（如 0.35 表示 35%），需要乘以 100
                    val_pct = val * 100 if abs(val) <= 1.0 else val
                    formatted_val = f"{val_pct:.2f}%"
                else:
                    formatted_val = f"{val:,.2f}"
            else:
                # 非数字类型（如日期、文本）直接转为字符串
                formatted_val = str(val)

            line_items.append(f"{col}：{formatted_val}")

        # 将这一期的所有字段用换行符拼接
        year_formatted_lines.append("\n".join(line_items))

    final_report = f"近五期年报主要财务数据如下：\n" + "\n\n".join(year_formatted_lines)

    last_type = df['report_type'].to_list()[0]

    if last_type != "年报":

        last_data = df[df.report_type == last_type].iloc[:2,:]

        # 假设你的原始数据存储在 df 中，使用 rename 修改列名
        # inplace=True 表示直接在原 DataFrame 上修改
        last_data.rename(columns=column_mapping, inplace=True)

        formatted_lines = []
        for _, row in last_data.iterrows():
            line_items = []
            for col, val in zip(last_data.columns, row):
                # 判断是否为数字类型
                if isinstance(val, (int, float)):
                    # 如果列名包含"率"，则格式化为保留两位小数并加上 %
                    if "率" in col:
                        # 如果原始数值本身就是小数（如 0.35 表示 35%），需要乘以 100
                        val_pct = val * 100 if abs(val) <= 1.0 else val
                        formatted_val = f"{val_pct:.2f}%"
                    else:
                        formatted_val = f"{val:,.2f}"
                else:
                    # 非数字类型（如日期、文本）直接转为字符串
                    formatted_val = str(val)

                line_items.append(f"{col}：{formatted_val}")

            # 将这一期的所有字段用换行符拼接
            formatted_lines.append("\n".join(line_items))

        # 3. 组装最终报告文本
        final_report = final_report + "\n\n" + f"近两期{last_type}主要财务数据如下：\n" + "\n\n".join(formatted_lines)

    return final_report

class MyOutputParser(StrOutputParser):
    def parse(self, text) -> str:
        return text.strip()


llm = ChatOpenAI(model="deepseek-v4-flash")
tools = [get_fina_info]
template = """
你是一位专业的股票基本面分析与投资研究专家。
你的任务是结合：
1. 实时获取的数据
2. 公司历史资料
3. 行业商业逻辑
4. 专业投资分析规则
对上市公司进行客观、严谨、深度的基本面诊断。
你的目标不是预测股价，而是分析公司的商业质量、竞争优势、成长逻辑以及潜在风险。
=============================
【信息来源优先级】
==============================
进行分析时，请遵循以下信息优先级：
第一优先级：
用户提供的数据、工具返回的数据、最新财务数据。
第二优先级：
公司事实资料（company_facts）。
第三优先级：
行业知识与商业逻辑（industry_knowledge）。
第四优先级：
通用分析规则（analysis_rules）。
如果不同来源之间存在冲突：
- 以最新、可验证的数据为准。
- 明确指出信息冲突。
- 不得自行编造缺失数据。
==============================
【决策与工具使用原则】
==============================
1. 按需调用工具：
根据用户问题类型，自主选择：
- 财务数据查询工具

不要为了调用工具而调用工具。
2. 数据驱动：
所有涉及以下内容：
- 收入
- 利润
- 毛利率
- 现金流
- 估值
- 市占率
- 产品规模
必须基于真实数据。
禁止：
- 编造财务数字
- 推测不存在的订单
- 将行业趋势当作公司事实
==============================
【知识库参考信息】
==============================
以下为公司事实资料：
{company_facts}
说明：
- 这些内容代表已经整理的公司历史信息。
- 可以用于理解公司的业务结构、产品、技术和历史事件。
- 不得擅自修改其中的数据。
------------------------------
以下为行业背景与商业逻辑：
{industry_knowledge}
说明：
- 用于判断行业空间、竞争格局、商业模式。
- 不代表当前公司的实际情况。
- 必须结合公司数据进行分析。
------------------------------
以下为基本面分析规则：
{analysis_rules}
说明：
- 这些规则用于辅助判断公司质量。
- 不应机械套用，需要结合行业特点分析。
==============================
【分析方法要求】
==============================
请按照以下逻辑进行分析：
一、数据与规则交叉验证
将公司的实际数据与分析规则进行匹配：
例如：
- 收入增长是否符合成长标准？
- 利润增速是否超过收入增速？
- 毛利率变化是否体现产品升级？
- 现金流是否支持利润真实性？
- 资产负债结构是否存在风险？
明确指出：
- 触发的积极信号
- 触发的风险信号
- 尚未验证的关键因素
------------------------------
二、商业模式分析
分析：
- 公司主要产品和收入来源
- 产品竞争壁垒
- 客户粘性
- 行业地位
- 上下游关系
重点回答：
“公司为什么能够赚钱？”
以及：
“这种赚钱能力是否可持续？”
------------------------------
三、成长逻辑分析
不要只关注收入增长。
需要分析：
- 新产品放量
- 市场空间
- 技术升级
- 国产替代
- 市占率变化
- 客户拓展
区分：
- 已经兑现的成长
- 预期中的成长
- 尚未验证的成长
------------------------------
四、风险评估
重点关注：
- 盈利质量下降
- 毛利率持续下滑
- 现金流恶化
- 资本开支压力
- 行业竞争加剧
- 技术路线变化
- 客户集中风险
- 安全生产风险
- 政策与国际环境风险

================================================================================
【 Telegram 消息专属输出排版规范 (CRITICAL FORMATTING RULES) 】
================================================================================

由于本诊断报告最终将在 Telegram 移动端客户端展示，为保证最佳的阅读体验与视觉质感，你必须严格遵循以下 Telegram 专属排版规则：

1. 禁用标准 Markdown 表格 (STRICTLY FORBIDDEN)：
   - 严禁使用 | 列分隔符 | 的 Markdown 表格语法！Telegram 移动端屏幕狭窄，多列表格会导致严重的强行折行与对齐错乱。
   - 所有多维财务数据、年度对比数据，必须拆解为【卡片式结构】或【树状分支结构】。

2. 树状分支与数据对齐 (Tree Structure)：
   - 涉及多项财务指标（如营收、归母、扣非）的平行对比时，使用 ├ 和 └ 字符构建清晰的树状层级：
     * 2025年
       ├ 营收：`25.85亿`（+9.93%） │ 归母：`3.20亿`（+18.00%）
       └ 扣非：`2.54亿`（+31.49%）
   - 单指标多年趋势，采用带箭头的平铺演进线表示：
     * 毛利率走势：`43.42%` (21) ➔ `41.16%` (24) ➔ `39.62%` (25)

3. 视觉高亮与数据卡片 (Inline Code Highlighting)：
   - 所有具体的【财务数值】、【百分比】、【比率】以及【核心定性结论】，必须用单反引号包围（即 Inline Code 格式，例如 `22.97%`、`中性偏正面`）。
   - 这会在 Telegram 界面中形成高亮灰色背景卡片，提升视觉重点的提炼效率。

4. 标题与符号锚点 (Visual Anchors)：
   - 每一个大章节标题前必须配有明确的 Emoji 视觉锚点（如 📊、💡、🔍、⚠️、🎯）。
   - 结论与诊断部分，优先采用带色彩倾向的标识符：
     * 正向/优质：✅ 或 🚀 或 `[良好]`
     * 中性/提示：📌 或 💡
     * 隐忧/警报：⚠️ 或 🔴 或 🟠

5. 控制行宽与分段 (Mobile Readability)：
   - 单段分析文字切忌过长，移动端每段尽量保持在 2-4 行以内。
   - 章节之间使用全划线 `---` 或空行进行视觉隔离，避免形成“文本大山”。

6. Telegram 模版输出参照：
   大纲结构示例：
   # 📊 [股票名称（代码）] 基本面深度诊断
   > **综合评级**：`[评级结果]`

   ---
   ### 💡 一、核心结论
   * **优质面**
     ├ **[维度]**：[简短分析]
     └ **[维度]**：[简短分析]

   ### 🔍 二、基本面深度诊断
   #### 1️⃣ 成长性：[一句话小结]
   * **[年份]**
     ├ 营收：`[数据]`（[增速]） │ 归母：`[数据]`（[增速]）
     └ 扣非：`[数据]`（[增速]）
   📌 **关键判断**：[分析文字]
"""
chat_prompt = ChatPromptTemplate.from_messages([
    ("system", template),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_tool_calling_agent(llm, tools, chat_prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

def run_fundamental_agent(user_input: str):

    # ======================
    # 多知识源检索
    # ======================
    rag_context = research_retrieve(
        user_input
    )

    # ======================
    # 注入Agent
    # ======================

    result = agent_executor.invoke(
        {
            "input": user_input,
            "company_facts":
                rag_context["company"],
            "industry_knowledge":
                rag_context["industry"],
            "analysis_rules":
                rag_context["rules"]
        }
    )


    return result["output"]

# ================= 新增：Telegram 机器人交互逻辑 =================
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """处理 /start 命令"""
    await update.message.reply_text("你好！我是你的股票分析助手，请直接发送类似“分析一下600699的基本面”的消息。")


async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """处理普通文本消息，调用你的 Agent，支持长消息分段与超时重试"""
    user_input = update.message.text
    await update.message.reply_text("🤔 正在分析，请稍候...")

    # 1. 获取 Agent 分析结果
    try:
        output = run_fundamental_agent(user_input)
    except Exception as e:
        output = f"分析出错：{str(e)}"

    # 2. 发送消息（带超时重试机制）
    max_retries = 5
    for attempt in range(1, max_retries + 1):
        try:
            # 如果消息长度在 Telegram 限制内，直接发送
            if len(output) <= 4096:
                await update.message.reply_text(output)
            else:
                # 消息超长，进行分段发送
                # 按换行符切分，尽量保持段落完整
                chunks = output.split('\n')
                current_chunk = ""

                for line in chunks:
                    # 如果单行本身就超过 4096 字符，强制按字符截断
                    if len(line) > 4096:
                        if current_chunk:
                            await update.message.reply_text(current_chunk)
                            current_chunk = ""
                        for i in range(0, len(line), 4096):
                            await update.message.reply_text(line[i:i + 4096])
                        continue

                    # 累加行数，如果加上当前行会超过限制，就先发送当前块
                    if len(current_chunk) + len(line) + 1 > 4096:
                        await update.message.reply_text(current_chunk)
                        current_chunk = line + "\n"
                    else:
                        current_chunk += line + "\n"

                # 发送最后剩余的内容
                if current_chunk:
                    await update.message.reply_text(current_chunk)

            # 如果发送成功，直接跳出重试循环
            break

        except TimedOut as e:
            # 捕获超时异常
            if attempt < max_retries:
                wait_time = attempt * 2  # 递增等待时间：2s, 4s, 6s...
                await update.message.reply_text(
                    f"⚠️ 消息发送超时，正在进行第 {attempt} 次重试，{wait_time}秒后继续..."
                )
                await asyncio.sleep(wait_time)
            else:
                await update.message.reply_text(
                    f"❌ 消息发送失败：连续 {max_retries} 次超时，请检查网络或稍后再试。"
                )
        except Exception as e:
            # 捕获其他非超时异常，直接报错退出，不进行重试
            await update.message.reply_text(f"❌ 发送消息时发生未知错误：{str(e)}")
            break

if __name__ == "__main__":
    # 从环境变量读取 Token
    token = os.getenv("TELEGRAM_BOT_TOKEN")
    if not token:
        raise ValueError("未在 .env 文件中找到 TELEGRAM_BOT_TOKEN")

    proxy_url = os.getenv(
        "TELEGRAM_PROXY"
    )

    builder = (
        ApplicationBuilder()
        .token(token)
    )

    if proxy_url:
        builder = (
            builder
            .proxy(proxy_url)
            .get_updates_proxy(proxy_url)
        )

    application = builder.build()

    application.add_handler(CommandHandler("start", start))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

    print("Telegram Bot 正在运行...")
    application.run_polling()