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

load_dotenv()


# ================= 你的原有逻辑（完全保留） =================
@tool
def get_fina_info(stock_code: str) -> str:
    """获取指定股票代码的财务数据。当需要获取财务数据的时使用此工具"""
    df = stock.finance.get_core_index(stock_code=stock_code)

    data = df[['report_date', 'report_type', 'total_rev', 'total_rev_yoy_gr', 'net_profit_attr_sh', 'net_profit_yoy_gr',
               'non_gaap_net_profit', 'non_gaap_net_profit_yoy_gr', 'roe_wtd', 'gross_margin', 'net_margin',
               'quick_ratio',
               'inv_turn_rate', 'acct_recv_turn_rate']]

    last_type = data['report_type'].to_list()[0]
    recent_df = data[data.report_type == last_type].iloc[:2, :].reset_index(drop=True)

    recent_df.rename(columns={"report_date": "报告日期", "report_type": "报告类型", "total_rev": "营业总收入",
                              "total_rev_yoy_gr": "营业总收入同比增长率", "net_profit_attr_sh": "归母净利润",
                              "net_profit_yoy_gr": "归母净利润同比增长率", "non_gaap_net_profit": "扣非净利润",
                              "non_gaap_net_profit_yoy_gr": "扣非净利润同比增长率", "roe_wtd": "加权净资产收益率",
                              "gross_margin": "销售毛利率", "net_margin": "销售净利率", "quick_ratio": "速动比率",
                              "inv_turn_rate": "存货周转率", "acct_recv_turn_rate": "应收账款周转率"
                              }, inplace=True)

    formatted_lines = []
    for _, row in recent_df.iterrows():
        line_items = []
        for col, val in zip(recent_df.columns, row):
            if isinstance(val, (int, float)):
                if "率" in col:
                    formatted_val = f"{val:.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))

    final_report = 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 = """
你是一个专业的股票分析助手。用户会输入一个股票代码，你可以使用工具来获取其财务数据，之后根据获取到的财务数据来分析这家公司的基本面情况。
"""
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)


# ================= 新增：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
    # 提示用户正在处理，因为 Agent 调用工具可能需要几秒钟
    await update.message.reply_text("🤔 正在分析，请稍候...")

    try:
        # 调用你原有的 Agent 逻辑
        res = agent_executor.invoke({"input": user_input})
        output = res["output"]
    except Exception as e:
        output = f"分析出错：{str(e)}"

    await update.message.reply_text(output)


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

    proxy_url = "http://127.0.0.1:8889"

    # 启动机器人
    application = ApplicationBuilder().token(token).proxy(proxy_url).get_updates_proxy(proxy_url).build()
    application.add_handler(CommandHandler("start", start))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

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