from typing import Annotated, TypedDict, Any
from langchain.agents import create_agent, AgentState
from langchain.chat_models import init_chat_model
from langchain.tools import tool, InjectedState
from langgraph.checkpoint.memory import InMemorySaver

from config.settings import LLM_API_KEY, LLM_BASE_URL, LLM_MODEL, VENT_REVIEW_PROMPT, VENT_CALCULATION_PROMPT, VENT_COMPARISON_PROMPT, INTENT_RECOGNITION_PROMPT
checkpointer = InMemorySaver()
# 导入所有配风审查工具
from tools.data_check_tools import (
    form_review,
    check_current_month_plan,
    check_data_by_gas_report,
    check_data_by_face_design,
    check_data_by_vent_report
)
from tools.clac_check_tools import (
    verify_coal_face_ventilation,
    verify_tunneling_face_ventilation,
    verify_chamber_ventilation,
    verify_other_points_ventilation
)
from tools.remote_clac_tools import (
    fetch_ventilation_data,
    recompute_ventilation,
    simulate_ventilation,
    fetch_field_translation
)

# 导入需风量对照与检查工具
from tools.vent_comparison_tools import (
    get_plan_required_wind,
    get_model_calculated_wind,
    get_sensor_wind,
    merge_wind_by_location,
    check_wind_compliance
)

# ============================================================
# SECTION 1: 自定义 Agent 状态
# ============================================================
class VentAgentState(AgentState):
    """配风审查 / 需风量对照 Agent 自定义状态"""
    vent_plan_data: dict[str, Any] | None          # 挂载配风计划完整JSON数据
    plan_wind_result: dict[str, Any] | None        # 工具1结果：配风计划需风量列表
    model_wind_result: dict[str, Any] | None       # 工具2结果：模型解算风量列表
    sensor_wind_result: dict[str, Any] | None      # 工具3结果：监测风量列表
    merged_wind_result: dict[str, Any] | None      # 工具4结果：合并后的风量对照数据

# ============================================================
# SECTION 2: 模型初始化
# ============================================================
model = init_chat_model(
    model=LLM_MODEL,
    api_key=LLM_API_KEY,
    base_url=LLM_BASE_URL,
    temperature=0,
    streaming=True,
    model_provider="openai",
    extra_body={"enable_thinking": False},
    timeout=300
)

# ============================================================
# SECTION 3: 工具列表
# ============================================================
review_tools = [
    form_review,
    check_current_month_plan,
    check_data_by_gas_report,
    check_data_by_face_design,
    check_data_by_vent_report,
    verify_coal_face_ventilation,
    verify_tunneling_face_ventilation,
    verify_chamber_ventilation,
    verify_other_points_ventilation,
]

calculation_tools = [
    fetch_ventilation_data,
    recompute_ventilation,
    simulate_ventilation,
    fetch_field_translation,
]

comparison_tools = [
    get_plan_required_wind,  # 需风量提取
    get_model_calculated_wind, # 三维模型解算风量提取
    get_sensor_wind,           # 巷道监测风量提取
    merge_wind_by_location,   # 地点风量智能合并
    check_wind_compliance, # 风量检查
]

# ============================================================
# SECTION 4: 创建 Agent 实例
# ============================================================
review_agent = create_agent(
    model=model,
    tools=review_tools,
    system_prompt=VENT_REVIEW_PROMPT,
    debug=False,
    checkpointer=checkpointer,
    state_schema=VentAgentState
)

calculation_agent = create_agent(
    model=model,
    tools=calculation_tools,
    system_prompt=VENT_CALCULATION_PROMPT,
    checkpointer=checkpointer,
    debug=False
)

comparison_agent = create_agent(
    model=model,
    tools=comparison_tools,
    system_prompt=VENT_COMPARISON_PROMPT,
    debug=False,
    checkpointer=checkpointer,
    state_schema=VentAgentState
)

intent_agent = create_agent(
    model=model,
    tools=[],
    system_prompt=INTENT_RECOGNITION_PROMPT,
    checkpointer=checkpointer,
    debug=False
)