"""第二阶段统一 Agent 内核。

完整调用顺序如下：

1. 根据入口选择客户顾问或运营助手人格；
2. 加载该人格的提示词、技能和允许使用的工具；
3. 将最近会话历史交给 LangChain Agent；
4. Agent 可自主进行多轮工具调用，包括知识检索和长期记忆；
5. 内核从工具执行记录中提取真实卡片、操作按钮和知识来源；
6. 最终返回稳定的 ``AgentReply``，前端不直接解析模型自由文本。

同一套内核承载两个 Agent 人格，但二者的身份、权限和工具集合相互隔离。
"""

import json
import re
from dataclasses import dataclass
from typing import Any, cast
from uuid import UUID, uuid4

from langchain.agents import create_agent
from langchain_core.messages import AIMessage
from langchain_core.runnables import RunnableConfig
from langsmith import Client, tracing_context

from zbt.core.config import Settings
from zbt.domains.identity.models import AdminUser, H5User
from zbt.harness.model_gateway import ModelGateway
from zbt.harness.policy import HarnessPolicyEngine
from zbt.harness.registries import PersonaRegistry, PromptRegistry, SkillRegistry
from zbt.harness.schemas import (
    AgentAction,
    AgentReply,
    AgentStructuredOutput,
    KnowledgeSourcesBlock,
    ProductRecommendationsBlock,
    ProductView,
    UiBlock,
)
from zbt.harness.tooling import ToolExecutionContext, ToolRegistry

CITATION_MARKER_PATTERN = re.compile(
    r"\s*<!--\s*ZBT_CITATIONS\s*:\s*(\[[\s\S]*?\])\s*-->\s*$",
    re.IGNORECASE,
)


@dataclass(frozen=True)
class ConversationTurn:
    """发送给模型的一条历史消息。"""

    role: str
    content: str


@dataclass(frozen=True)
class AgentInvocation:
    """一次 Agent 调用需要的全部上下文。"""

    persona: str
    message: str
    h5_user: H5User | None = None
    admin_user: AdminUser | None = None
    history: tuple[ConversationTurn, ...] = ()


class AgentKernel:
    """负责模型编排，不直接实现投保、检索或分析等具体业务。"""

    def __init__(
        self,
        *,
        settings: Settings,
        model_gateway: ModelGateway,
        tools: ToolRegistry,
        prompts: PromptRegistry | None = None,
        skills: SkillRegistry | None = None,
        personas: PersonaRegistry | None = None,
        policy: HarnessPolicyEngine | None = None,
    ) -> None:
        self._settings = settings
        self._model_gateway = model_gateway
        self._tools = tools
        self._prompts = prompts or PromptRegistry()
        self._skills = skills or SkillRegistry()
        self._personas = personas or PersonaRegistry()
        self._policy = policy or HarnessPolicyEngine()
        self._langsmith_client = (
            Client(
                api_url=settings.langsmith_endpoint,
                api_key=settings.langsmith_api_key,
            )
            if settings.langsmith_tracing and settings.langsmith_api_key
            else None
        )

    def reply(self, invocation: AgentInvocation) -> AgentReply:
        """运行一次 Agent，并将模型输出与真实工具结果组装成前端协议。"""

        persona = self._personas.get(invocation.persona)
        system_prompt = self._compose_prompt(persona.prompt_id, persona.skill_ids)
        context = ToolExecutionContext(
            persona=persona.id,
            h5_user=invocation.h5_user,
            admin_user=invocation.admin_user,
        )
        # 这里只构建当前人格允许使用的工具，模型看不到其他人格的能力。
        tools = self._tools.build(
            names=persona.tool_names,
            context=context,
            policy_engine=self._policy,
        )
        agent = create_agent(
            model=self._model_gateway.chat_model(persona.id),
            tools=tools,
            system_prompt=system_prompt,
            name=f"zhibaotong-{persona.id}",
        )
        trace_id = uuid4()
        # 限制历史窗口，避免会话无限增长导致请求变慢和上下文成本失控。
        messages = [
            {"role": turn.role, "content": turn.content} for turn in invocation.history[-12:]
        ]
        messages.append({"role": "user", "content": invocation.message})
        with tracing_context(
            project_name=self._settings.langsmith_project,
            enabled=self._langsmith_client is not None,
            client=self._langsmith_client,
            tags=["zhibaotong", "s2", persona.id],
            metadata={
                "persona": persona.id,
                "channel": ("customer_h5" if persona.id == "customer" else "operation_admin"),
                "principal_id": context.principal_id,
            },
        ):
            run_config: RunnableConfig = {"run_id": trace_id}
            result = agent.invoke(
                cast(Any, {"messages": messages}),
                config=run_config,
            )
        structured = self._structured_output(
            cast(dict[str, Any], result),
            context,
        )
        blocks: list[UiBlock] = [
            block
            for execution in context.executions
            for block in execution.result.blocks
            if block.type != "knowledge_sources"
        ]
        # 知识来源只能从本次真实工具执行记录中选择，不能由模型凭空生成。
        knowledge_sources = self._knowledge_sources_block(
            context,
            structured.cited_knowledge_chunk_ids,
        )
        if knowledge_sources is not None:
            blocks.append(knowledge_sources)
        recommendation = self._recommendation_block(
            context,
            structured.selected_product_codes,
        )
        if recommendation is not None:
            blocks.insert(0, recommendation)
        actions = [
            action for execution in context.executions for action in execution.result.actions
        ]
        if structured.suggested_action is not None:
            actions.append(structured.suggested_action)
        unique_actions = _deduplicate_actions(actions)
        return AgentReply(
            text=structured.message,
            cards=[block.model_dump(mode="json") for block in blocks],
            trace_id=str(trace_id),
            trace_url=self._trace_url(trace_id),
            actions=[action.model_dump(mode="json") for action in unique_actions],
            invoked_tools=tuple(execution.name for execution in context.executions),
        )

    def _compose_prompt(
        self,
        prompt_id: str,
        skill_ids: tuple[str, ...],
    ) -> str:
        prompt = self._prompts.get(prompt_id)
        skill_sections = []
        for skill_id in skill_ids:
            skill = self._skills.get(skill_id)
            skill_sections.append(f"## Skill：{skill.name}\n\n{skill.instructions}")
        return f"{prompt}\n\n# 已加载技能\n\n" + "\n\n".join(skill_sections)

    @staticmethod
    def _structured_output(
        result: dict[str, Any],
        context: ToolExecutionContext,
    ) -> AgentStructuredOutput:
        structured_response = result.get("structured_response")
        if isinstance(structured_response, AgentStructuredOutput):
            return structured_response
        if isinstance(structured_response, dict):
            return AgentStructuredOutput.model_validate(structured_response)
        messages = result.get("messages", [])
        final = messages[-1] if messages else None
        text = (
            str(final.content)
            if isinstance(final, AIMessage)
            else "暂时无法生成有效回答，请稍后重试。"
        )
        text, cited_chunk_ids = _extract_citation_marker(text)
        selected_codes: list[str] = []
        for execution in context.executions:
            if execution.name != "list_available_products":
                continue
            raw_items = execution.result.data.get("items", [])
            if not isinstance(raw_items, list):
                continue
            for item in raw_items:
                if not isinstance(item, dict):
                    continue
                code = str(item.get("product_code", ""))
                name = str(item.get("name", ""))
                if code and (code in text or name in text):
                    selected_codes.append(code)
        return AgentStructuredOutput(
            message=text,
            selected_product_codes=list(dict.fromkeys(selected_codes))[:2],
            cited_knowledge_chunk_ids=cited_chunk_ids,
        )

    @staticmethod
    def _knowledge_sources_block(
        context: ToolExecutionContext,
        cited_chunk_ids: list[str],
    ) -> KnowledgeSourcesBlock | None:
        """根据模型声明使用的片段编号，生成聚合后的知识来源卡片。"""

        if not cited_chunk_ids:
            return None
        candidates: dict[str, dict[str, Any]] = {}
        for execution in context.executions:
            if execution.name != "search_insurance_knowledge":
                continue
            raw_items = execution.result.data.get("items", [])
            if not isinstance(raw_items, list):
                continue
            for item in raw_items:
                if not isinstance(item, dict):
                    continue
                chunk_id = str(item.get("chunk_id", ""))
                if chunk_id:
                    candidates[chunk_id] = item
        selected = [
            candidates[chunk_id]
            for chunk_id in dict.fromkeys(cited_chunk_ids)
            if chunk_id in candidates
        ][:5]
        if not selected:
            return None
        # 同一文档同一版本的多个片段聚合成一个来源，避免前端重复展示。
        grouped: dict[tuple[str, int], list[dict[str, Any]]] = {}
        for item in selected:
            key = (
                str(item["document_id"]),
                int(item["document_version"]),
            )
            evidence = grouped.setdefault(key, [])
            content = str(item["content"])[:180]
            if (
                len(evidence) < 2
                and all(str(existing["content"])[:180] != content for existing in evidence)
            ):
                evidence.append(item)
        return KnowledgeSourcesBlock(
            items=[
                {
                    "document_id": document_id,
                    "document_version": document_version,
                    "title": str(evidence[0]["title"]),
                    "source_name": str(evidence[0]["source_name"]),
                    "excerpt": str(evidence[0]["content"])[:180],
                    "excerpts": [
                        str(item["content"])[:180]
                        for item in evidence
                    ],
                    "score": max(float(item["score"]) for item in evidence),
                }
                for (document_id, document_version), evidence in list(
                    grouped.items()
                )[:5]
                if evidence
            ]
        )

    @staticmethod
    def _recommendation_block(
        context: ToolExecutionContext,
        selected_codes: list[str],
    ) -> ProductRecommendationsBlock | None:
        if not selected_codes:
            return None
        product_items: list[dict[str, Any]] = []
        for execution in context.executions:
            if execution.name != "list_available_products":
                continue
            raw_items = execution.result.data.get("items", [])
            if isinstance(raw_items, list):
                product_items.extend(item for item in raw_items if isinstance(item, dict))
        selected = [
            item
            for code in selected_codes
            for item in product_items
            if item.get("product_code") == code
        ][:2]
        if not selected:
            return None
        views = [
            ProductView.model_validate(
                {
                    "product_id": item["product_id"],
                    "product_code": item["product_code"],
                    "name": item["name"],
                    "category": item["category"],
                    "summary": item["summary"],
                    "product_version_id": item["product_version_id"],
                    "plans": item["plans"],
                }
            )
            for item in selected
        ]
        return ProductRecommendationsBlock(items=views)

    def _trace_url(self, trace_id: UUID) -> str | None:
        if self._langsmith_client is None:
            return None
        try:
            run = self._langsmith_client.read_run(trace_id)
            return self._langsmith_client.get_run_url(
                run=run,
                project_name=self._settings.langsmith_project,
            )
        except Exception:
            # Trace 上传可能有短暂延迟，运行详情接口会再次尝试解析。
            return None

    def trace_url(self, trace_id: str) -> str | None:
        try:
            return self._trace_url(UUID(trace_id))
        except ValueError:
            return None


def _deduplicate_actions(actions: list[AgentAction]) -> list[AgentAction]:
    unique: list[AgentAction] = []
    seen: set[tuple[str, str]] = set()
    for action in actions:
        key = (action.type, str(sorted(action.payload.items())))
        if key not in seen:
            seen.add(key)
            unique.append(action)
    return unique


def _extract_citation_marker(text: str) -> tuple[str, list[str]]:
    match = CITATION_MARKER_PATTERN.search(text)
    if match is None:
        return text.strip(), []
    cleaned = text[: match.start()].rstrip()
    try:
        raw_ids = json.loads(match.group(1))
    except json.JSONDecodeError:
        return cleaned, []
    if not isinstance(raw_ids, list):
        return cleaned, []
    cited_ids = [
        item
        for item in raw_ids
        if isinstance(item, str) and item.strip()
    ]
    return cleaned, list(dict.fromkeys(cited_ids))[:5]
