from typing import Any, List, Sequence, Union

import langchain_core
from langchain.tools import BaseTool
from langchain_core.prompts import PromptTemplate,ChatPromptTemplate,SystemMessagePromptTemplate,MessagesPlaceholder,HumanMessagePromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.agents import AgentExecutor, Agent, create_tool_calling_agent,create_openai_functions_agent,create_structured_chat_agent
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
from langchain.agents.format_scratchpad.openai_tools import (
    format_to_openai_tool_messages,
)
from langchain import hub


from src.agent.tool_rate import RegionRateTool,RankingRateTool
from src.agent.tool_monitor import MonitorPointTool

# def create_rate_agent(llm, tools: List[BaseTool],prompt: PromptTemplate = None, 
#     tools_renderer: ToolsRenderer = render_text_description_and_args,
#     verbose: bool = False,**args):
#     missing_vars = {"tools", "tool_names", "agent_scratchpad"}.difference(
#         prompt.input_variables + list(prompt.partial_variables)
#     )
#     if missing_vars:
#         raise ValueError(f"Prompt missing required variables: {missing_vars}")

#     prompt = prompt.partial(
#         tools=tools_renderer(list(tools)),
#         tool_names=", ".join([t.name for t in tools]),
#     )
#     if stop_sequence:
#         stop = ["\nObservation"] if stop_sequence is True else stop_sequence
#         llm_with_stop = llm.bind(stop=stop)
#     else:
#         llm_with_stop = llm

#     agent = (
#         RunnablePassthrough.assign(
#             agent_scratchpad=lambda x: format_log_to_str(x["intermediate_steps"]),
#         )
#         | prompt
#         | llm_with_stop
#     )
#     return agent


class RateAgent:
    def __init__(self, llm, tools: List[BaseTool],prompt: PromptTemplate = None, verbose: bool = False,**args):
        # if not prompt:
        #     raise ValueError("PromptTemplate is required")

        prompt = hub.pull("hwchase17/openai-tools-agent")
        agent = create_tool_calling_agent(llm, tools, prompt)

        self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=verbose)    
        
    def exec(self, prompt_args: dict = {}, stream: bool = False):
        return self.agent_executor.invoke(input=prompt_args)

    def stream(self, prompt_args: dict = {}):
        for step in self.agent_executor.stream(prompt_args):
            yield step


# 适配 structured_chat_agent 的 prompt
ONLINE_RATE_SYSTEM_PROMPT = """你是一个专门处理地质监测点信息及监测设备在线率分析的AI助手。你可以通过调用专门的工具来分析和展示不同维度的在线率数据。

你需要：
1. 理解用户意图，将用户问题映射到合适的分析类型
2. 确保必要参数完整，如果缺少参数则提示用户缺少参数
3. 如果参数完整，则调用相应的分析工具获取数据
4. 生成清晰的分析报告，包括数据解读
5. 工具返回的数据务必用 markdown 格式的数据表格进行展示
6. 对异常情况（如数据缺失、参数错误）提供友好的解释和建议，不要自己创造虚拟数据

注意事项：
- 时间格式统一使用：YYYY-MM-DD
- 地区名称需要包含行政级别（如：福建省、厦门市）
- 数据展示优先使用 markdown 格式的数据表格，保证数据的完整性，并配合文字说明
- 百分比数据保留两位小数

您可以使用以下工具：

{tools}

使用 JSON 对象指定工具，提供一个 action 键（工具名称）和一个 action_input 键（工具输入） 。

有效的 "action" 值： "Final Answer" 或 {tool_names}

每个 $JSON_BLOB 只提供一个操作，如下所示：
    ``` 
    {{
    "action": $TOOL_NAME,
    "action_input": $INPUT,
    }}
    ```

按照以下格式：

Question: 输入要回答的问题
Thought: 考虑前后步骤
Action:
    ```
    $JSON_BLOB
    ```
Observation: 操作结果
...（重复 Thought/Action/Observation N 次）
Thought: 我知道如何回复
Action:
    ```
    {{
    "action": "Final Answer",
    "action_input": "最终回复给人类",
    }}
    ```
开始！始终以有效的单个操作的 JSON 对象回复。如有必要，请使用工具。如果你知道答案，请直接回复。
你的回复格式为 Action:```$JSON_BLOB```然后 Observation。
"""

PROMPT_AGENT_HUMAN = """{input}\n\n {agent_scratchpad}\n (请注意，无论如何都要以 JSON 对象回复。工具返回的数据必须使用表格展示，包含在最终输出中，并且要保证数据的完整性)"""
PROMPT_AGENT_SYS_VARS = [ "tool_names", "tools"]

class RateAgentV2:
    def __init__(self, llm, tools: List[BaseTool],prompt: PromptTemplate = None, verbose: bool = False,**args):
        prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template(ONLINE_RATE_SYSTEM_PROMPT),
            MessagesPlaceholder(variable_name="chat_history", optional=True),
            HumanMessagePromptTemplate.from_template(PROMPT_AGENT_HUMAN)
        ])

        agent = create_structured_chat_agent(llm, tools, prompt)

        self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=verbose, return_intermediate_steps=True)    
        
    def exec(self, prompt_args: dict = {}, stream: bool = False):
        return self.agent_executor.invoke(input=prompt_args)

    def stream(self, prompt_args: dict = {}):
        for step in self.agent_executor.stream(prompt_args):
            yield step

def new_rate_agent(llm, verbose: bool = False,**args):

    if args['tool_base_url']:
        tool_base_url = args['tool_base_url']
    else:
        tool_base_url = const_base_url

    tools = [
        RegionRateTool(base_url=tool_base_url),
        RankingRateTool(base_url=tool_base_url),
        MonitorPointTool(base_url=tool_base_url)
    ]

    # 使用 LangChain 的工具调用代理
    agent = RateAgentV2(llm=llm, tools=tools, verbose=verbose, **args)
    return agent