"""Agent-execution subgraph for multi-step tool-enabled execution."""

from __future__ import annotations

import json
import logging
import os
from typing import Any

from langchain_core.tools import BaseTool
from langgraph.graph import END, START, StateGraph

from app.a2a_server import create_text_part, get_current_updater
from a2a.types import TaskState
from app.server import get_llm_provider_registry, get_mcp_client
from config.settings import LLMRole
from graph.state import RoomAgentGraphState
from llm_json_parse import JsonParserWithRepair

from .state import AgentExecutionState, FinalOutput, PlannerStep


logger = logging.getLogger(__name__)
DEFAULT_STEP_LIMIT = 6
USER_VISIBLE_UNFINISHED_MESSAGE = "任务暂时无法完成, 请稍后重试。"
PLANNER_OUTPUT_SCHEMA = {
    "type": "object",
    "properties": {
        "step_type": {
            "type": "string",
            "enum": ["reason", "toolcall", "final_output"],
        },
        "is_done": {"type": "boolean"},
        "reason_summary": {"type": "string"},
        "tool_name": {"type": "string"},
        "tool_args": {"type": "object"},
        "final_output": {
            "type": "object",
            "properties": {
                "message": {"type": "string", "minLength": 1},
                "summary": {"type": ["string", "null"]},
                "metadata": {"type": ["object", "null"]},
            },
            "required": ["message"],
            "additionalProperties": False,
        },
    },
    "required": ["step_type", "is_done"],
    "additionalProperties": False,
}


async def agent_execution(state: RoomAgentGraphState) -> RoomAgentGraphState:
    """Run the internal agent-execution subgraph and return an outer-state patch."""
    result = await compile_agent_execution_subgraph().ainvoke(
        {
            "user_input": state.get("user_input", ""),
            "conversation_text": state.get("conversation_text", ""),
            "intent": dict(state.get("intent", {})),
            "selected_tools": list(state.get("selected_tools", [])),
            "metadata": dict(state.get("metadata", {})),
        }
    )
    return result.get("outer_state_patch", {})


def compile_agent_execution_subgraph() -> Any:
    """Compile the internal agent-execution subgraph."""
    graph = StateGraph(AgentExecutionState)
    graph.add_node("subgraph_input_transform", subgraph_input_transform)
    graph.add_node("agent_plan_step", agent_plan_step)
    graph.add_node("agent_dispatch_step", agent_dispatch_step)
    graph.add_node("agent_record_reason", agent_record_reason)
    graph.add_node("agent_execute_toolcall", agent_execute_toolcall)
    graph.add_node("agent_finalize_output", agent_finalize_output)
    graph.add_node("agent_handle_limit_or_error", agent_handle_limit_or_error)
    graph.add_node("agent_submit_reply", agent_submit_reply)
    graph.add_node("subgraph_output_transform", subgraph_output_transform)

    graph.add_edge(START, "subgraph_input_transform")
    graph.add_edge("subgraph_input_transform", "agent_plan_step")
    graph.add_conditional_edges(
        "agent_plan_step",
        route_after_plan_step,
        {
            "dispatch": "agent_dispatch_step",
            "handle_error": "agent_handle_limit_or_error",
        },
    )
    graph.add_conditional_edges(
        "agent_dispatch_step",
        route_after_dispatch_step,
        {
            "reason": "agent_record_reason",
            "toolcall": "agent_execute_toolcall",
            "final_output": "agent_finalize_output",
            "handle_error": "agent_handle_limit_or_error",
        },
    )
    graph.add_edge("agent_record_reason", "agent_plan_step")
    graph.add_conditional_edges(
        "agent_execute_toolcall",
        route_after_toolcall,
        {
            "plan": "agent_plan_step",
            "handle_error": "agent_handle_limit_or_error",
        },
    )
    graph.add_edge("agent_finalize_output", "agent_submit_reply")
    graph.add_edge("agent_handle_limit_or_error", "agent_submit_reply")
    graph.add_edge("agent_submit_reply", "subgraph_output_transform")
    graph.add_edge("subgraph_output_transform", END)
    return graph.compile()


async def subgraph_input_transform(state: AgentExecutionState) -> AgentExecutionState:
    """Initialize internal state for the agent loop."""
    available_tools, tool_instances = await _load_selected_tools(state.get("selected_tools", []))
    return {
        "user_input": state.get("user_input", "").strip(),
        "conversation_text": state.get("conversation_text", "").strip()
        or state.get("user_input", "").strip(),
        "intent": dict(state.get("intent", {})),
        "metadata": dict(state.get("metadata", {})),
        "available_tools": available_tools,
        "available_tool_instances": tool_instances,
        "step_count": 0,
        "step_limit": DEFAULT_STEP_LIMIT,
        "step_history": [],
        "raw_model_outputs": [],
        "tool_results": [],
        "replan_used": False,
    }


async def agent_plan_step(state: AgentExecutionState) -> AgentExecutionState:
    """Plan the next agent step with the powerful model."""
    if state.get("step_count", 0) >= state.get("step_limit", DEFAULT_STEP_LIMIT):
        return {
            "loop_decision": "handle_error",
            "terminal_error": {
                "type": "agent_step_limit_exceeded",
                "message": (
                    f"Agent step limit exceeded before completion "
                    f"(limit={state.get('step_limit', DEFAULT_STEP_LIMIT)})."
                ),
                "source_node": "agent_plan_step",
                "retryable": False,
            },
        }

    provider = _get_powerful_provider()
    raw_output = await provider.complete_text(
        _build_planner_messages(state),
        temperature=0,
        json_mode=True,
    )
    parsed = await JsonParserWithRepair(llm_provider=provider)(
        raw_output,
        schema=PLANNER_OUTPUT_SCHEMA,
    )
    current_step = _normalize_planner_step(parsed)
    _validate_planner_step(current_step)

    raw_model_outputs = list(state.get("raw_model_outputs", []))
    raw_model_outputs.append(raw_output)

    return {
        "current_step": current_step,
        "raw_model_outputs": raw_model_outputs,
        "step_count": state.get("step_count", 0) + 1,
        "loop_decision": "dispatch",
    }


def agent_dispatch_step(state: AgentExecutionState) -> AgentExecutionState:
    """No-op node used to keep dispatch logic explicit in the graph."""
    current_step = state.get("current_step", {})
    if not current_step:
        return {
            "loop_decision": "handle_error",
            "terminal_error": {
                "type": "agent_dispatch_error",
                "message": "Planner did not produce a current_step.",
                "source_node": "agent_dispatch_step",
                "retryable": False,
            },
        }
    return {}


def agent_record_reason(state: AgentExecutionState) -> AgentExecutionState:
    """Record a reasoning-only step and continue the loop."""
    current_step = state.get("current_step", {})
    step_history = list(state.get("step_history", []))

    updater = None
    try:
        updater = get_current_updater()
    except RuntimeError:
        logger.debug("No TaskUpdater available; skip reply submit from agent_execution subgraph.")

    if updater and current_step.get("reason_summary"):
        message = updater.new_agent_message(
            parts=[create_text_part(current_step.get("reason_summary", "思考中...").strip())],
        )
        _ = updater.update_status(TaskState.working, message)

    step_history.append(
        {
            "step_index": state.get("step_count", 0),
            "step_type": "reason",
            "reason_summary": current_step.get("reason_summary", "").strip(),
        }
    )
    return {"step_history": step_history}


async def agent_execute_toolcall(state: AgentExecutionState) -> AgentExecutionState:
    """Execute a single tool call and either continue or fail."""
    current_step = state.get("current_step", {})
    tool_name = str(current_step.get("tool_name", "")).strip()
    tool_args = current_step.get("tool_args", {}) or {}
    if _is_dry_run_enabled():
        observation = _build_dry_run_observation(tool_name=tool_name, tool_args=tool_args)
        logger.info(
            "RA_DRY_RUN intercepted MCP tool call tool=%s observation=%s",
            tool_name,
            json.dumps(observation, ensure_ascii=False, default=str),
        )
        return _build_tool_success_state(
            state,
            tool_name=tool_name,
            tool_args=tool_args,
            observation=observation,
        )

    updater = None
    try:
        updater = get_current_updater()
    except RuntimeError:
        logger.debug("No TaskUpdater available; skip reply submit from agent_execution subgraph.")

    tool_instances = state.get("available_tool_instances", {})
    tool = tool_instances.get(tool_name)
    if tool is None:
        return _build_tool_failure_state(
            state,
            tool_name=tool_name,
            tool_args=tool_args,
            error_message=f"Tool '{tool_name}' is not available.",
            source_node="agent_execute_toolcall",
        )

    try:
        observation = await _invoke_tool(tool, tool_args)
    except Exception as exc:
        return _build_tool_failure_state(
            state,
            tool_name=tool_name,
            tool_args=tool_args,
            error_message=f"{type(exc).__name__}: {exc}",
            source_node="agent_execute_toolcall",
        )

    if updater:
        message = updater.new_agent_message(parts=[create_text_part(f"已执行 {tool_name}")])
        await updater.update_status(TaskState.working, message)

    return _build_tool_success_state(
        state,
        tool_name=tool_name,
        tool_args=tool_args,
        observation=observation,
    )


def agent_finalize_output(state: AgentExecutionState) -> AgentExecutionState:
    """Persist the final output and prepare to exit the subgraph."""
    current_step = state.get("current_step", {})
    final_output = current_step.get("final_output", {})
    step_history = list(state.get("step_history", []))
    step_history.append(
        {
            "step_index": state.get("step_count", 0),
            "step_type": "final_output",
            "message": final_output.get("message", ""),
            "summary": final_output.get("summary"),
        }
    )
    return {
        "final_output": final_output,
        "step_history": step_history,
    }


def agent_handle_limit_or_error(state: AgentExecutionState) -> AgentExecutionState:
    """Build terminal error state for an unfinished agent run."""
    terminal_error = dict(state.get("terminal_error", {}))
    if not terminal_error:
        terminal_error = {
            "type": "agent_execution_error",
            "message": "Agent execution failed without a structured error.",
            "source_node": "agent_handle_limit_or_error",
            "retryable": False,
        }
    return {"terminal_error": terminal_error}


async def agent_submit_reply(state: AgentExecutionState) -> AgentExecutionState:
    """Submit a user-visible reply through the active A2A updater when available."""
    reply = _resolve_terminal_reply_text(state)
    if not reply:
        return {}

    try:
        updater = get_current_updater()
    except RuntimeError:
        logger.debug("No TaskUpdater available; skip reply submit from agent_execution subgraph.")
        return {}

    message = updater.new_agent_message(parts=[create_text_part(reply)])
    await updater.complete(message)
    return {}


def subgraph_output_transform(state: AgentExecutionState) -> AgentExecutionState:
    """Map the internal terminal state back to the outer graph state."""
    tool_call_history = _build_tool_call_history(state.get("tool_results", []))
    final_output = _normalize_final_output(state.get("final_output"))
    terminal_error = state.get("terminal_error")

    if final_output and not terminal_error:
        execution_result = {
            "type": "agent_final_output",
            "message": final_output["message"],
            "tool_call_history": tool_call_history,
        }
        if final_output.get("summary") is not None:
            execution_result["summary"] = final_output.get("summary")
        if final_output.get("metadata") is not None:
            execution_result["metadata"] = final_output.get("metadata")
        outer_patch: RoomAgentGraphState = {
            "status": "completed",
            "next_action": "agent_execution",
            "execution_result": execution_result,
            "tool_call_history": tool_call_history,
        }
        return {"outer_state_patch": outer_patch}

    error_message = (
        str(terminal_error.get("message", "")).strip()
        if isinstance(terminal_error, dict)
        else "Agent execution did not complete."
    )
    outer_patch = {
        "status": "failed",
        "next_action": "agent_execution",
        "error": terminal_error
        or {
            "type": "agent_execution_error",
            "message": error_message or "Agent execution did not complete.",
            "source_node": "subgraph_output_transform",
            "retryable": False,
        },
        "execution_result": {
            "type": "agent_execution_unfinished",
            "message": USER_VISIBLE_UNFINISHED_MESSAGE,
            "tool_call_history": tool_call_history,
            "unfinished": True,
        },
        "tool_call_history": tool_call_history,
    }
    return {"outer_state_patch": outer_patch}


def _resolve_terminal_reply_text(state: AgentExecutionState) -> str:
    final_output = _normalize_final_output(state.get("final_output"))
    message = str(final_output.get("message", "")).strip()
    if message and not state.get("terminal_error"):
        return message
    return USER_VISIBLE_UNFINISHED_MESSAGE


def route_after_plan_step(state: AgentExecutionState) -> str:
    """Route after planning."""
    return state.get("loop_decision", "handle_error")


def route_after_dispatch_step(state: AgentExecutionState) -> str:
    """Dispatch the current planner step."""
    if state.get("loop_decision") == "handle_error":
        return "handle_error"
    current_step = state.get("current_step", {})
    return str(current_step.get("step_type", "handle_error"))


def route_after_toolcall(state: AgentExecutionState) -> str:
    """Route after a tool call finishes."""
    return state.get("loop_decision", "handle_error")


def _normalize_planner_step(parsed: dict[str, Any]) -> PlannerStep:
    step: PlannerStep = {
        "step_type": str(parsed["step_type"]),
        "is_done": bool(parsed["is_done"]),
    }
    if "reason_summary" in parsed:
        step["reason_summary"] = str(parsed.get("reason_summary") or "").strip()
    if "tool_name" in parsed:
        step["tool_name"] = str(parsed.get("tool_name") or "").strip()
    if "tool_args" in parsed and isinstance(parsed["tool_args"], dict):
        step["tool_args"] = parsed["tool_args"]
    if "final_output" in parsed and isinstance(parsed["final_output"], dict):
        step["final_output"] = _normalize_final_output(parsed["final_output"])
    return step


def _normalize_final_output(payload: dict[str, Any] | FinalOutput | None) -> FinalOutput:
    if not payload:
        return {}
    normalized: FinalOutput = {"message": str(payload.get("message", "")).strip()}
    if payload.get("summary") is not None:
        normalized["summary"] = str(payload.get("summary"))
    if payload.get("metadata") is not None and isinstance(payload.get("metadata"), dict):
        normalized["metadata"] = payload.get("metadata")
    return normalized


def _validate_planner_step(step: PlannerStep) -> None:
    step_type = step.get("step_type")
    is_done = bool(step.get("is_done"))
    if step_type == "reason":
        if is_done:
            raise ValueError("Planner step_type=reason must have is_done=false.")
        if not step.get("reason_summary"):
            raise ValueError("Planner step_type=reason requires reason_summary.")
        return
    if step_type == "toolcall":
        if is_done:
            raise ValueError("Planner step_type=toolcall must have is_done=false.")
        if not step.get("tool_name"):
            raise ValueError("Planner step_type=toolcall requires tool_name.")
        if not isinstance(step.get("tool_args", {}), dict):
            raise ValueError("Planner step_type=toolcall requires tool_args object.")
        return
    if step_type == "final_output":
        if not is_done:
            raise ValueError("Planner step_type=final_output must have is_done=true.")
        final_output = step.get("final_output", {})
        if not final_output.get("message"):
            raise ValueError("Planner step_type=final_output requires final_output.message.")
        return
    raise ValueError(f"Unsupported planner step_type={step_type!r}.")


def _render_tool_list(tools: list[dict[str, Any]]) -> str:
    if not tools:
        return "(none)"

    rendered_tools = []
    for tool in tools:
        rendered_tools.append(
            "\n".join(
                [
                    f"- {tool.get('name', '<unknown>')}: {tool.get('description') or 'No description'}",
                    f"  args_schema: {json.dumps(tool.get('args_schema', {}), ensure_ascii=False, sort_keys=True)}",
                ]
            )
        )
    return "\n".join(rendered_tools)


def _build_planner_messages(state: AgentExecutionState) -> list[dict[str, str]]:
    return [
        {
            "role": "system",
            "content": (
                "你是一个房间的智能管理助手。你的工作是管理房间的各种设备和事务, 根据用户的需求进行推理和决策。\n"
                "\n"
                "**你每轮的任务很简单：决定下一步该做什么。** 你有三个选择：\n"
                "\n"
                "1. **思考与分析**：如果问题比较复杂, 你可以先思考一下思路, 理清思路后再执行动作。简洁地说出你的想法即可。\n"
                "\n"
                "2. **执行工具**：当你知道需要做什么时, 就拿起一个可用的工具去执行。记住每次只能用一个工具。告诉系统工具的名称和所需参数。\n"
                "   你必须严格参考每个工具给出的 args_schema 来构造 tool_args。"
                "不要编造 schema 里没有的字段, 字段名和结构必须与 schema 对齐。\n"
                "\n"
                "3. **给出最终答案**：当任务完成或问题解决了, 就告诉用户结果。这是唯一一个真正结束对话的方式。\n"
                "\n"
                "**重要规则**: \n"
                "• 你最多可以进行 6 个步骤。如果无法完成, 就坦诚告诉用户。\n"
                "• 思考和工具调用都会让步骤计数器增加, 所以要高效地思考。\n"
                "• 如果工具执行失败, 你可以重新规划一次。再次失败的话就放弃吧。\n"
                "• 你必须只输出一个 JSON 对象, 不要输出额外解释、Markdown 或自然语言前后缀。\n"
                "• JSON 字段必须与系统要求一致：step_type、is_done, 以及按需提供 reason_summary、tool_name、tool_args、final_output。\n"
                "\n"
                "输出要求示例：\n"
                '{"step_type":"reason","is_done":false,"reason_summary":"先确认当前状态"}\n'
                '{"step_type":"toolcall","is_done":false,"tool_name":"climate_get_state","tool_args":{"entity_id":"climate.lab"}}\n'
                '{"step_type":"final_output","is_done":true,"final_output":{"message":"实验室空调当前正在运行"}}'
            ),
        },
        {
            "role": "user",
            "content": (
                f"用户的原始输入: {state.get('user_input', '')}\n"
                f"对话上下文: {state.get('conversation_text', '')}\n"
                f"识别的意图: {state.get('intent', {})}\n"
                f"可用工具: {_render_tool_list(state.get('available_tools', []))}\n"
                f"当前步数: {state.get('step_count', 0)} / 最多步数: {state.get('step_limit', DEFAULT_STEP_LIMIT)}\n"
                f"执行历史: {state.get('step_history', [])}\n"
                f"工具结果: {state.get('tool_results', [])}\n"
                f"已重新规划: {state.get('replan_used', False)}\n"
                "请只返回一个 JSON 对象。"
            ),
        },
    ]


def _get_powerful_provider() -> Any:
    provider = get_llm_provider_registry().get(LLMRole.POWERFUL)
    if provider is None:
        raise RuntimeError(f"LLM provider is unavailable for role={LLMRole.POWERFUL.value}")
    return provider


async def _load_selected_tools(
    selected_tools: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], dict[str, BaseTool]]:
    client = get_mcp_client()
    selected_names = {
        str(tool.get("name", "")).strip()
        for tool in selected_tools
        if isinstance(tool, dict) and str(tool.get("name", "")).strip()
    }
    if client is None:
        return list(selected_tools), {}

    tools = await client.get_tools()
    available_tools = []
    tool_instances: dict[str, BaseTool] = {}
    for tool in tools:
        if selected_names and tool.name not in selected_names:
            continue
        tool_instances[tool.name] = tool
        available_tools.append(
            {
                "name": tool.name,
                "description": tool.description or "",
                "args_schema": _extract_tool_schema(tool),
            }
        )
    if available_tools:
        return available_tools, tool_instances
    return list(selected_tools), {}


async def _invoke_tool(tool: BaseTool, tool_args: dict[str, Any]) -> Any:
    if hasattr(tool, "ainvoke"):
        return await tool.ainvoke(tool_args)
    if hasattr(tool, "arun"):
        return await tool.arun(tool_args)
    if hasattr(tool, "invoke"):
        return tool.invoke(tool_args)
    if hasattr(tool, "run"):
        return tool.run(tool_args)
    raise RuntimeError(f"Tool {getattr(tool, 'name', '<unknown>')} is not invokable.")


def _extract_tool_schema(tool: BaseTool) -> dict[str, Any]:
    get_input_schema = getattr(tool, "get_input_schema", None)
    if callable(get_input_schema):
        schema_model = get_input_schema()
        model_json_schema = getattr(schema_model, "model_json_schema", None)
        if callable(model_json_schema):
            schema = model_json_schema()
            return schema if isinstance(schema, dict) else {}

    args_schema = getattr(tool, "args", {}) or {}
    if args_schema:
        return {
            "type": "object",
            "properties": args_schema,
        }

    return {}


def _build_tool_failure_state(
    state: AgentExecutionState,
    *,
    tool_name: str,
    tool_args: dict[str, Any],
    error_message: str,
    source_node: str,
) -> AgentExecutionState:
    tool_results = list(state.get("tool_results", []))
    step_history = list(state.get("step_history", []))
    entry = {
        "step_index": state.get("step_count", 0),
        "tool_name": tool_name,
        "tool_args": tool_args,
        "error": error_message,
        "args_summary": _summarize_value(tool_args),
        "error_summary": _summarize_value(error_message),
    }
    tool_results.append(entry)
    step_history.append(
        {
            "step_index": state.get("step_count", 0),
            "step_type": "toolcall",
            "tool_name": tool_name,
            "args_summary": entry["args_summary"],
            "error_summary": entry["error_summary"],
        }
    )
    if not state.get("replan_used", False):
        return {
            "tool_results": tool_results,
            "step_history": step_history,
            "replan_used": True,
            "loop_decision": "plan",
        }

    return {
        "tool_results": tool_results,
        "step_history": step_history,
        "loop_decision": "handle_error",
        "terminal_error": {
            "type": "tool_execution_error",
            "message": error_message,
            "source_node": source_node,
            "retryable": False,
        },
    }


def _build_tool_success_state(
    state: AgentExecutionState,
    *,
    tool_name: str,
    tool_args: dict[str, Any],
    observation: Any,
) -> AgentExecutionState:
    tool_results = list(state.get("tool_results", []))
    step_history = list(state.get("step_history", []))
    result_summary = _summarize_value(observation)
    tool_results.append(
        {
            "step_index": state.get("step_count", 0),
            "tool_name": tool_name,
            "tool_args": tool_args,
            "observation": observation,
            "args_summary": _summarize_value(tool_args),
            "result_summary": result_summary,
        }
    )
    step_history.append(
        {
            "step_index": state.get("step_count", 0),
            "step_type": "toolcall",
            "tool_name": tool_name,
            "args_summary": _summarize_value(tool_args),
            "result_summary": result_summary,
        }
    )
    return {
        "tool_results": tool_results,
        "step_history": step_history,
        "loop_decision": "plan",
    }


def _build_tool_call_history(tool_results: list[dict[str, Any]]) -> list[dict[str, Any]]:
    history = []
    for entry in tool_results:
        item = {
            "step_index": entry.get("step_index"),
            "tool_name": entry.get("tool_name"),
            "args_summary": entry.get("args_summary", ""),
        }
        if entry.get("result_summary"):
            item["result_summary"] = entry["result_summary"]
        if entry.get("error_summary"):
            item["error_summary"] = entry["error_summary"]
        history.append(item)
    return history


def _summarize_value(value: Any, *, limit: int = 180) -> str:
    if isinstance(value, str):
        text = value
    else:
        try:
            text = json.dumps(value, ensure_ascii=False, default=str)
        except TypeError:
            text = str(value)
    text = text.strip()
    if len(text) <= limit:
        return text
    return text[: limit - 3] + "..."


def _is_dry_run_enabled() -> bool:
    value = os.getenv("RA_DRY_RUN", "").strip().lower()
    return value not in {"", "0", "false", "no", "off"}


def _build_dry_run_observation(*, tool_name: str, tool_args: dict[str, Any]) -> dict[str, Any]:
    return {
        "dry_run": True,
        "tool_name": tool_name,
        "tool_args": tool_args,
        "message": "RA_DRY_RUN enabled; skipped real MCP invocation.",
    }
