from platform import java_ver
from langgraph.graph import StateGraph
from langgraph.graph.state import RunnableConfig
from pydantic import BaseModel
import requests
from typing import Annotated, Type, TypedDict, Optional
# from langchain_core.annotations import InjectedToolArg

from langchain_openai import ChatOpenAI

from alayaflow.common.config import settings

from langgraph.graph import START, END
from langchain.tools import tool

"""
结点参数
- 每个结点使用的时候有固定的参数
    - 1. 工作流定义中传入
- 包含结点的工作流创建的时候有全局变量，来配置工作流
    - 2. 工作流创建时输入的全局变量
- 结点运行时有输入参数
    - 3. 工作流运行时输入 （来自用户）
    - 4. 运行时上游结点的输出

调用方式
- 放入langgraph中作单独结点
- 放入langflow component实现中
- ~放入LLM结点作tools~
"""

# Core

def search(query, url, item_cnt, api_key):
    # return requests.get(
    #     f"{url}/search",
    #     headers={"Authorization": f"Bearer {search_api_key}"},
    #     params={"query": query, "item": return_item_cnt}
    # ).json()
    
    # Fake run
    return f"Get request from {url} with headers={{'Authorization': 'Bearer {api_key}'}} and params={{'query': {query}, 'item': {item_cnt}}}"

# @tool
# def search_as_tool(
#     query: str, 
#     # 告诉 LangChain：这些参数不要给 LLM 看，由代码注入
#     config: Annotated[dict, InjectedToolArg] 
# ):
#     """搜索工具，用于获取特定信息。"""
#     # 从注入的 config 中提取配置
#     return search(
#         query=query,
#         url=config["url"],
#         api_key=config["api_key"],
#         item_cnt=config.get("item_cnt", 3)
#     )


# Workflow developer layer

class MyState(BaseModel):
    query: str
    search_result: Optional[str] = None

# class SearchNode:
#     def __init__(self, key: str):
#         self.key = key
    
#     def __call__(self, state: MyState):
#         new_state = state.model_copy()
#         new_state.search_result = search(
#             state.query,
#             "https://xxx",    # 1
#             3,       # 1
#             self.key,  # 2
#         )
#         return new_state


def create_search_node(api_key: str):
    
    def search_node(state: MyState, config: RunnableConfig):
        result = search(
            query=state.query, 
            url=config["configurable"]["search_url"],
            api_key=config["configurable"]["search_api_key"],
            item_cnt=3,
        )
        return {"search_result": result}
    return search_node


def create_graph(search_api_key):
    builder = StateGraph(MyState)

    # 添加节点
    builder.add_node("search", create_search_node(search_api_key))

    # 定义边
    builder.add_edge(START, "search")
    builder.add_edge("search", END)

    return builder.compile()

class WFConfig(TypedDict):
    search_api_key: str
    search_url: str

def get_config_schema() -> Type[TypedDict]:
    return WFConfig


if __name__ == "__main__":
    # Example usage
    graph = create_graph(search_api_key="your-api-key-here")
    result = graph.invoke({
        "query": "search query",
    }, config={
        "configurable": WFConfig(
            search_api_key="your-api-key-here",
            search_url="https://your-search-api-url",
        )
    })
    print(f"result: {result}")

