"""离线示例：共享术语表服务、检索工具和模型调用前的上下文补充。"""
from contextlib import asynccontextmanager

import config
from extensions import define_tool, get_service


def enabled() -> bool:
    return config.get_bool("CONTEXT_DEMO_ENABLED", False)


@asynccontextmanager
async def glossary_service(dependencies):
    # 实际接入时在这里延迟导入客户端、打开连接，并在 finally 中关闭。
    glossary = {"风量": "单位时间内通过某一断面的空气体积。", "风速": "空气流动的速度。"}
    try:
        yield glossary
    finally:
        glossary.clear()


@define_tool(group="扩展示例", label="查询术语")
def lookup_demo_term(term: str) -> str:
    """查询离线演示术语表，仅用于演示模块接入。"""
    return get_service("demo_glossary").get(term, "术语表中没有该词条")


def glossary_middleware(services):
    # 模块关闭时不会加载可选依赖。工厂只组装中间件，不创建连接。
    from langchain.agents.middleware.types import AgentMiddleware
    class DemoGlossaryMiddleware(AgentMiddleware):
        def enrich(self, request):
            message = request.system_message
            if message is None:
                return request
            terms = "、".join(sorted(services["demo_glossary"]))
            text = f"\n演示术语表可查询：{terms}。需要时调用 lookup_demo_term。"
            content = (message.content + text if isinstance(message.content, str)
                       else [*message.content, {"type": "text", "text": text}])
            return request.override(system_message=message.model_copy(update={
                "content": content,
            }))

        def wrap_model_call(self, request, handler):
            return handler(self.enrich(request))

        async def awrap_model_call(self, request, handler):
            return await handler(self.enrich(request))

    return DemoGlossaryMiddleware()


def setup(registry):
    registry.register_service("demo_glossary", glossary_service)
    registry.register_tool(lookup_demo_term, requires=["demo_glossary"])
    registry.register_middleware("DemoGlossaryMiddleware", glossary_middleware,
                                requires=["demo_glossary"], targets=["main"])