import operator

from typing import TypedDict, List, Annotated
from langgraph.graph import END, StateGraph
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from src.multi_agent.config import LLM
from src.multi_agent.util.weather_search import search_weather

llm = LLM


class WeatherState(TypedDict):
    user_request: str
    city: str
    result: str


def requirement_analysis(state: WeatherState):
    user_request = state.get('user_request')
    analysis_prompt = ChatPromptTemplate([
        ('system', '你是一个天气查询agent，你负责从用户的需求中提取出需要查询天气的城市或地区，你只需要输出这个地名，不需要输出其他任何内容'),
        ('user', f'当前的用户需求是:'),
        MessagesPlaceholder(variable_name="user_request"),
    ])
    analysis_chain = analysis_prompt | llm
    llm_resp = analysis_chain.invoke({'user_request': [("user", user_request)]})
    return {'city': llm_resp.content}


def get_weather_data(state: WeatherState):
    city = state.get('city')
    weather_data = search_weather(city)
    return {'result': weather_data}


def enter_chain(input_data: dict):
    return {'user_request': input_data['next_plan']}


workflow = StateGraph(WeatherState)
# 其实完全没必要使用子agent，但演示一下子agent如何并网

workflow.add_node('需求分析', requirement_analysis)
workflow.add_node('获取天气', get_weather_data)

workflow.set_entry_point('需求分析')

workflow.add_edge('需求分析', '获取天气')
workflow.add_edge('获取天气', END)

weather_graph = workflow.compile()
weather_chain = enter_chain | weather_graph


if __name__ == '__main__':
    requirement_analysis({})