from __future__ import annotations

import asyncio
import contextlib
import json
import os
import time
from collections import deque
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit
from typing import Any, AsyncGenerator, Callable, Literal, Mapping
from uuid import uuid4

import httpx
from pydantic import BaseModel, Field

from app.config import settings

try:
    from agents import (
        Agent,
        ItemHelpers,
        ModelResponse,
        MultiProvider,
        OpenAIProvider,
        OpenAIResponsesModel,
        RunConfig,
        RunContextWrapper,
        Runner,
        SQLiteSession,
        Usage,
        function_tool,
        trace,
    )
    from agents.models.interface import ModelProvider
    from agents.models.multi_provider import MultiProviderMap
    from agents.models.openai_responses import Converter as ResponsesConverter
    from agents.util._json import _to_dump_compatible
except ImportError:  # pragma: no cover - 依赖缺失时走降级分支
    Agent = None
    ItemHelpers = None
    ModelResponse = None
    MultiProvider = None
    OpenAIProvider = None
    OpenAIResponsesModel = None
    RunConfig = None
    RunContextWrapper = Any
    Runner = None
    SQLiteSession = None
    Usage = None
    function_tool = None
    trace = None
    MultiProviderMap = None
    ModelProvider = object
    ResponsesConverter = None
    _to_dump_compatible = None


def sdk_available() -> bool:
    return all(
        dependency is not None
        for dependency in (
            Agent,
            ModelResponse,
            OpenAIProvider,
            OpenAIResponsesModel,
            ResponsesConverter,
            RunConfig,
            Runner,
            SQLiteSession,
            Usage,
            function_tool,
            trace,
        )
    )


def has_configured_openai_api_key() -> bool:
    value = get_effective_openai_api_key()
    if not value:
        return False
    return "please-set-your-key" not in value


def _is_placeholder_openai_key(value: str | None) -> bool:
    return "please-set-your-key" in str(value or "").strip()


@lru_cache(maxsize=1)
def load_codex_auth_openai_key() -> str:
    auth_path = Path.home() / ".codex" / "auth.json"
    try:
        payload = json.loads(auth_path.read_text("utf-8"))
    except Exception:
        return ""
    value = str(payload.get("OPENAI_API_KEY", "")).strip() if isinstance(payload, dict) else ""
    if not value or _is_placeholder_openai_key(value):
        return ""
    return value


def get_effective_openai_api_key(env: Mapping[str, str] | None = None) -> str:
    env_map = env if env is not None else os.environ
    env_value = str(env_map.get("OPENAI_API_KEY", "")).strip()
    settings_value = settings.openai_api_key.strip()
    auth_value = load_codex_auth_openai_key()

    if env_value and not _is_placeholder_openai_key(env_value):
        return env_value
    if settings_value and not _is_placeholder_openai_key(settings_value):
        return settings_value
    if auth_value and not _is_placeholder_openai_key(auth_value):
        return auth_value
    return env_value or settings_value or auth_value


@lru_cache(maxsize=8)
def load_available_openai_models(
    base_url: str,
    api_key: str,
) -> tuple[str, ...]:
    normalized_base_url = base_url.rstrip("/")
    if not normalized_base_url or not api_key.strip():
        return ()

    try:
        with httpx.Client(timeout=10.0) as client:
            response = client.get(
                f"{normalized_base_url}/models",
                headers={"Authorization": f"Bearer {api_key.strip()}"},
            )
        response.raise_for_status()
        payload = response.json()
    except Exception:
        return ()

    data = payload.get("data") if isinstance(payload, dict) else None
    if not isinstance(data, list):
        return ()

    ids: list[str] = []
    for item in data:
        if not isinstance(item, dict):
            continue
        model_id = item.get("id")
        if isinstance(model_id, str) and model_id.strip():
            ids.append(model_id.strip())
    return tuple(ids)


@dataclass(frozen=True)
class DocumentAiModelSpec:
    key: str
    title: str
    description: str
    gateway_model: str
    resolved_combo: str | None = None
    resolved_runtime_model: str | None = None
    status: str = "active"
    is_default: bool = False


@dataclass(frozen=True)
class DocumentAiProfileSpec:
    id: str
    title: str
    description: str
    instructions: str
    session_mode: Literal["off", "page", "workspace"] = "page"
    tool_names: tuple[str, ...] = ()
    is_default: bool = False
    status: str = "active"


@dataclass(frozen=True)
class DocumentAiToolSpec:
    name: str
    title: str
    description: str
    scope: Literal["document", "tree", "workspace"] = "document"
    mode: Literal["read", "write"] = "read"
    status: str = "active"
    version: str = "v1"


DEFAULT_OMNIROUTE_BASE_URL = "http://localhost:20128/v1"


DOCUMENT_AI_MODEL_SPECS: tuple[DocumentAiModelSpec, ...] = (
    DocumentAiModelSpec(
        key="gpt-5.4",
        title="GPT-5.4",
        description="文档页默认主写模型键，经由 OmniRoute 二次路由。",
        gateway_model="gpt-5.4",
        resolved_combo="slow",
        resolved_runtime_model="gpt-5.4",
        is_default=True,
    ),
    DocumentAiModelSpec(
        key="gpt-5.4-mini",
        title="GPT-5.4 Mini",
        description="更偏快响应的轻量模型键，经由 OmniRoute 二次路由。",
        gateway_model="gpt-5.4-mini",
        resolved_runtime_model="gpt-5.4-mini",
    ),
    DocumentAiModelSpec(
        key="gpt-5.3-codex",
        title="GPT-5.3 Codex",
        description="偏代码与结构化执行的模型键，经由 OmniRoute 二次路由。",
        gateway_model="gpt-5.3-codex",
        resolved_combo="codex",
        resolved_runtime_model="gpt-5.3-codex",
    ),
    DocumentAiModelSpec(
        key="gemini-3.1-pro-preview",
        title="Gemini 3.1 Pro Preview",
        description="偏长上下文与深推理的模型键，经由 OmniRoute 二次路由。",
        gateway_model="gemini-3.1-pro-preview",
        resolved_runtime_model="gemini-3.1-pro-preview",
    ),
    DocumentAiModelSpec(
        key="third",
        title="Third",
        description="第三路由保底模型键，经由 OmniRoute combo 直通。",
        gateway_model="third",
        resolved_combo="third",
        resolved_runtime_model="third",
    ),
)


DOCUMENT_AI_TOOL_SPECS: tuple[DocumentAiToolSpec, ...] = (
    DocumentAiToolSpec(
        name="doc_get",
        title="读取当前页摘要",
        description="读取当前页块摘要，适合先理解页面结构。",
        mode="read",
    ),
    DocumentAiToolSpec(
        name="doc_find",
        title="按块查找",
        description="按文本查找当前页块，适合定位待改写内容。",
        mode="read",
    ),
    DocumentAiToolSpec(
        name="doc_insert_blocks",
        title="插入块",
        description="向当前页插入新块，适合新增标题和段落。",
        mode="write",
    ),
    DocumentAiToolSpec(
        name="doc_replace_range",
        title="改写块文本",
        description="改写当前页指定块文本。",
        mode="write",
    ),
    DocumentAiToolSpec(
        name="slash_run",
        title="斜杠命令",
        description="执行当前页允许的斜杠命令；当前仅限重命名当前页。",
        mode="write",
    ),
)


DOCUMENT_AI_PROFILE_SPECS: tuple[DocumentAiProfileSpec, ...] = (
    DocumentAiProfileSpec(
        id="page_writer_ai_first",
        title="AI 主写",
        description="AI 优先直接进入主编辑区起草，人负责审核和继续编辑。",
        instructions=(
            "当前 profile：AI 优先直接进入主编辑区起草，人负责审核和继续编辑。"
            " 如果用户目标明确，优先输出可直接落入当前页的修改结果，而不是长篇空谈。"
        ),
        tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
        is_default=True,
    ),
    DocumentAiProfileSpec(
        id="page_writer_strict",
        title="严格审慎",
        description="更保守地读取、定位、改写，优先减少误写。",
        instructions=(
            "当前 profile：严格审慎。写入前优先读取和定位目标块；若信息不足，先解释缺口，不要冒进写入。"
        ),
        tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
    ),
    DocumentAiProfileSpec(
        id="page_writer_polish",
        title="润色整理",
        description="更偏改写、收敛、结构整理，不主动扩写超出当前页语境的内容。",
        instructions=(
            "当前 profile：润色整理。优先压缩、润色、重组当前页已有内容，不主动扩写超出当前页证据范围的新事实。"
        ),
        tool_names=tuple(spec.name for spec in DOCUMENT_AI_TOOL_SPECS),
    ),
)


def _document_ai_model_spec_by_key(model_key: str | None) -> DocumentAiModelSpec | None:
    value = str(model_key or "").strip()
    if not value:
        return None
    return next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.key == value), None)


def _document_ai_model_spec_by_gateway_model(gateway_model: str | None) -> DocumentAiModelSpec | None:
    value = str(gateway_model or "").strip()
    if not value:
        return None
    return next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.gateway_model == value), None)


def _resolve_default_document_ai_model_key(
    env: Mapping[str, str] | None = None,
) -> str:
    env_map = env if env is not None else os.environ
    for value in _iter_candidate_model_names(
        request_model=None,
        configured_default_model=settings.mnote_ai_default_model,
        env=env_map,
    ):
        spec = _document_ai_model_spec_by_key(value) or _document_ai_model_spec_by_gateway_model(value)
        if spec is not None:
            return spec.key
    default_spec = next((spec for spec in DOCUMENT_AI_MODEL_SPECS if spec.is_default), DOCUMENT_AI_MODEL_SPECS[0])
    return default_spec.key


def _resolve_default_document_ai_profile_id() -> str:
    default_spec = next((spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.is_default), DOCUMENT_AI_PROFILE_SPECS[0])
    return default_spec.id


def resolve_document_ai_profile(profile_id: str | None) -> DocumentAiProfileSpec:
    value = str(profile_id or "").strip()
    if value:
        picked = next((spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.id == value), None)
        if picked is not None:
            return picked
    default_id = _resolve_default_document_ai_profile_id()
    return next(spec for spec in DOCUMENT_AI_PROFILE_SPECS if spec.id == default_id)


def build_document_ai_config_payload(
    env: Mapping[str, str] | None = None,
) -> dict[str, Any]:
    env_map = env if env is not None else os.environ
    default_model_key = _resolve_default_document_ai_model_key(env_map)
    default_profile_id = _resolve_default_document_ai_profile_id()
    raw_base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip() or DEFAULT_OMNIROUTE_BASE_URL
    return {
        "provider": "online",
        "transport": "openai_responses",
        "baseUrl": normalize_loopback_base_url(raw_base_url),
        "sessionEnabled": True,
        "defaultModelKey": default_model_key,
        "defaultProfileId": default_profile_id,
        "models": [
            {
                "key": spec.key,
                "title": spec.title,
                "description": spec.description,
                "gatewayModel": spec.gateway_model,
                "resolvedCombo": spec.resolved_combo,
                "resolvedRuntimeModel": spec.resolved_runtime_model,
                "status": spec.status,
                "default": spec.key == default_model_key,
            }
            for spec in DOCUMENT_AI_MODEL_SPECS
        ],
        "profiles": [
            {
                "id": spec.id,
                "title": spec.title,
                "description": spec.description,
                "sessionMode": spec.session_mode,
                "toolNames": list(spec.tool_names),
                "status": spec.status,
                "default": spec.id == default_profile_id,
            }
            for spec in DOCUMENT_AI_PROFILE_SPECS
        ],
        "tools": [
            {
                "name": spec.name,
                "title": spec.title,
                "description": spec.description,
                "scope": spec.scope,
                "mode": spec.mode,
                "status": spec.status,
                "version": spec.version,
            }
            for spec in DOCUMENT_AI_TOOL_SPECS
        ],
    }


def normalize_loopback_base_url(base_url: str) -> str:
    value = str(base_url or "").strip()
    if not value:
        return DEFAULT_OMNIROUTE_BASE_URL
    parsed = urlsplit(value)
    if parsed.hostname != "127.0.0.1":
        return value
    netloc = parsed.netloc.replace("127.0.0.1", "localhost", 1)
    return urlunsplit((parsed.scheme, netloc, parsed.path, parsed.query, parsed.fragment))


def _iter_candidate_model_names(
    *,
    request_model: str | None,
    configured_default_model: str | None,
    env: Mapping[str, str],
) -> tuple[str, ...]:
    values = (
        (request_model or "").strip(),
        (configured_default_model or "").strip(),
        str(env.get("ONLINE_AI_MODEL", "")).strip(),
        str(env.get("OPENAI_MODEL", "")).strip(),
        str(env.get("OPENAI_DEFAULT_MODEL", "")).strip(),
        str(env.get("CLAW_DEFAULT_MODEL", "")).strip(),
    )
    return tuple(value for value in values if value)


def _extract_model_prefix(model_name: str) -> str | None:
    value = str(model_name or "").strip()
    if "/" not in value:
        return None
    prefix, _ = value.split("/", 1)
    return prefix.strip() or None


class GatewayPrefixedOpenAIProvider(ModelProvider):
    def __init__(
        self,
        prefix: str,
        *,
        api_key: str | None,
        base_url: str | None,
        use_responses: bool = True,
    ) -> None:
        self._prefix = prefix
        self._api_key = api_key or ""
        self._base_url = base_url or ""
        self._use_responses = use_responses

    def get_model(self, model_name: str | None):
        full_model_name = f"{self._prefix}/{model_name}" if model_name else self._prefix
        return CompatOpenAIResponsesModel(
            model=full_model_name,
            openai_client=None,
            model_is_explicit=True,
            base_url=self._base_url,
            api_key=self._api_key,
            use_responses=self._use_responses,
        )


class DocumentAiModelProvider(ModelProvider):
    def __init__(
        self,
        *,
        provider_map: Mapping[str, ModelProvider],
        api_key: str | None,
        base_url: str | None,
        use_responses: bool = True,
    ) -> None:
        self._provider_map = dict(provider_map)
        self._api_key = api_key or ""
        self._base_url = base_url or ""
        self._use_responses = use_responses

    def get_model(self, model_name: str | None):
        value = str(model_name or "").strip()
        if "/" in value:
            prefix, suffix = value.split("/", 1)
            provider = self._provider_map.get(prefix)
            if provider is None:
                provider = GatewayPrefixedOpenAIProvider(
                    prefix,
                    api_key=self._api_key,
                    base_url=self._base_url,
                    use_responses=self._use_responses,
                )
            return provider.get_model(suffix)

        return CompatOpenAIResponsesModel(
            model=value or None,
            openai_client=None,
            model_is_explicit=bool(value),
            base_url=self._base_url,
            api_key=self._api_key,
            use_responses=self._use_responses,
        )


def build_agent_usage_from_response_usage(response_usage: Any):
    if response_usage is None:
        return Usage()

    def pick(name: str, default: Any = None) -> Any:
        if isinstance(response_usage, Mapping):
            return response_usage.get(name, default)
        return getattr(response_usage, name, default)

    input_tokens = int(pick("input_tokens", 0) or 0)
    output_tokens = int(pick("output_tokens", 0) or 0)
    total_tokens_raw = pick("total_tokens", None)
    total_tokens = input_tokens + output_tokens
    if total_tokens_raw not in (None, ""):
        try:
            total_tokens = int(total_tokens_raw)
        except (TypeError, ValueError):
            total_tokens = input_tokens + output_tokens
    return Usage(
        requests=1,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=total_tokens,
        input_tokens_details=pick("input_tokens_details", None),
        output_tokens_details=pick("output_tokens_details", None),
    )


def _is_openai_omit(value: Any) -> bool:
    return value is None or value.__class__.__name__ == "Omit"


def parse_gateway_raw_response_text(raw_text: str) -> dict[str, Any]:
    text = raw_text.strip()
    if not text:
        raise RuntimeError("网关未返回有效响应体")

    if text.startswith("{"):
        payload = json.loads(text)
        if not isinstance(payload, dict):
            raise RuntimeError("网关响应不是对象")
        return payload

    parsed_payload: dict[str, Any] | None = None
    for block in text.split("\n\n"):
        lines = [
            line[len("data:") :].strip()
            for line in block.splitlines()
            if line.startswith("data:")
        ]
        if not lines:
            continue
        try:
            data = json.loads("\n".join(lines))
        except json.JSONDecodeError:
            continue
        if isinstance(data, dict) and isinstance(data.get("response"), dict):
            parsed_payload = data["response"]
            continue
        if isinstance(data, dict):
            parsed_payload = data

    if parsed_payload is None:
        raise RuntimeError(f"无法解析网关 Responses 返回：{text[:200]}")
    return parsed_payload


class CompatOpenAIResponsesModel(OpenAIResponsesModel):
    def __init__(
        self,
        model,
        openai_client,
        *,
        model_is_explicit: bool = True,
        base_url: str,
        api_key: str,
        use_responses: bool = True,
    ) -> None:
        super().__init__(model=model, openai_client=openai_client, model_is_explicit=model_is_explicit)
        self._gateway_base_url = base_url.rstrip("/")
        self._gateway_api_key = api_key
        self._gateway_use_responses = use_responses

    async def get_response(
        self,
        system_instructions,
        input,
        model_settings,
        tools,
        output_schema,
        handoffs,
        tracing,
        previous_response_id: str | None = None,
        conversation_id: str | None = None,
        prompt=None,
    ):
        list_input = ItemHelpers.input_to_new_input_list(input)
        list_input = _to_dump_compatible(list_input)
        list_input = self._remove_openai_responses_api_incompatible_fields(list_input)

        if model_settings.parallel_tool_calls and tools:
            parallel_tool_calls = True
        elif model_settings.parallel_tool_calls is False:
            parallel_tool_calls = False
        else:
            parallel_tool_calls = None

        tool_choice = ResponsesConverter.convert_tool_choice(model_settings.tool_choice)
        converted_tools = ResponsesConverter.convert_tools(tools, handoffs)
        converted_tools_payload = _to_dump_compatible(converted_tools.tools)
        response_format = ResponsesConverter.get_response_format(output_schema)
        should_omit_model = prompt is not None and not self._model_is_explicit
        should_omit_tools = prompt is not None and len(converted_tools_payload) == 0

        include_set: set[str] = set(converted_tools.includes)
        if model_settings.response_include is not None:
            include_set.update(model_settings.response_include)
        if model_settings.top_logprobs is not None:
            include_set.add("message.output_text.logprobs")
        include = list(include_set)

        payload: dict[str, Any] = {
            "previous_response_id": previous_response_id,
            "conversation": conversation_id,
            "instructions": system_instructions,
            "model": None if should_omit_model else self.model,
            "input": list_input,
            "include": include or None,
            "tools": None if should_omit_tools else converted_tools_payload,
            "prompt": prompt,
            "temperature": model_settings.temperature,
            "top_p": model_settings.top_p,
            "truncation": model_settings.truncation,
            "max_output_tokens": model_settings.max_tokens,
            "tool_choice": tool_choice,
            "parallel_tool_calls": parallel_tool_calls,
            "text": response_format,
            "store": model_settings.store,
            "prompt_cache_retention": model_settings.prompt_cache_retention,
            "reasoning": model_settings.reasoning,
            "metadata": model_settings.metadata,
        }
        if model_settings.top_logprobs is not None:
            payload["top_logprobs"] = model_settings.top_logprobs
        if model_settings.verbosity is not None:
            if isinstance(payload.get("text"), dict):
                payload["text"]["verbosity"] = model_settings.verbosity
            else:
                payload["text"] = {"verbosity": model_settings.verbosity}
        if model_settings.extra_body:
            payload.update(dict(model_settings.extra_body))

        payload = {
            key: value
            for key, value in payload.items()
            if not _is_openai_omit(value)
        }
        headers = {
            "Authorization": f"Bearer {self._gateway_api_key}",
            "Content-Type": "application/json",
        }
        if isinstance(model_settings.extra_headers, Mapping):
            headers.update({str(k): str(v) for k, v in model_settings.extra_headers.items()})

        async with httpx.AsyncClient(timeout=60.0) as client:
            raw_response = await client.post(
                f"{self._gateway_base_url}/responses",
                headers=headers,
                json=payload,
                params=model_settings.extra_query,
            )
        raw_response.raise_for_status()
        parsed_payload = parse_gateway_raw_response_text(raw_response.text)

        if isinstance(parsed_payload.get("error"), dict):
            raise RuntimeError(
                str(parsed_payload["error"].get("message") or "Responses API 返回错误")
            )

        return ModelResponse(
            output=parsed_payload.get("output") or [],
            usage=build_agent_usage_from_response_usage(parsed_payload.get("usage")),
            response_id=str(parsed_payload.get("id") or "") or None,
        )


def wrap_openai_model_for_gateway_compat(model: Any):
    if OpenAIResponsesModel is not None and isinstance(model, OpenAIResponsesModel):
        return CompatOpenAIResponsesModel(
            model=model.model,
            openai_client=model._client,
            model_is_explicit=getattr(model, "_model_is_explicit", True),
            base_url=getattr(getattr(model, "_client", None), "base_url", ""),
            api_key="",
        )
    return model


def build_document_model_provider(
    *,
    request_model: str | None,
    configured_default_model: str | None,
    env: Mapping[str, str] | None = None,
    available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
):
    if not sdk_available():
        raise RuntimeError("openai-agents 未安装，无法构建文档页 model provider")

    env_map = env if env is not None else os.environ
    base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip()
    api_key = get_effective_openai_api_key(env_map)
    loader = available_models_loader or load_available_openai_models
    available_models = list(loader(base_url, api_key)) if base_url and api_key else []

    prefixes: set[str] = set()
    for model_name in _iter_candidate_model_names(
        request_model=request_model,
        configured_default_model=configured_default_model,
        env=env_map,
    ):
        prefix = _extract_model_prefix(model_name)
        if prefix:
            prefixes.add(prefix)
    for model_name in available_models:
        prefix = _extract_model_prefix(str(model_name))
        if prefix:
            prefixes.add(prefix)
    provider_map: dict[str, ModelProvider] = {}
    for prefix in sorted(prefixes):
        if prefix in {"openai", "litellm"}:
            continue
        provider_map[prefix] = GatewayPrefixedOpenAIProvider(
            prefix,
            api_key=api_key,
            base_url=base_url,
            use_responses=True,
        )

    return DocumentAiModelProvider(
        provider_map=provider_map,
        api_key=api_key,
        base_url=base_url,
        use_responses=True,
    )


def resolve_document_agent_model(
    *,
    request_model: str | None,
    configured_default_model: str | None,
    env: Mapping[str, str] | None = None,
    available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
) -> str:
    env_map = env if env is not None else os.environ
    candidates = _iter_candidate_model_names(
        request_model=request_model,
        configured_default_model=configured_default_model,
        env=env_map,
    )
    raw_model = candidates[0] if candidates else ""

    if not raw_model:
        raise RuntimeError(
            "未配置文档页 AI model；请设置 MNOTE_AI_DEFAULT_MODEL 或在请求中显式传入 provider/model"
        )

    if "/" in raw_model:
        return raw_model

    base_url = str(env_map.get("OPENAI_BASE_URL", "")).strip()
    api_key = get_effective_openai_api_key(env_map)
    if not base_url or not api_key:
        raise RuntimeError(
            f"文档页 AI model '{raw_model}' 缺少 provider 前缀，且当前无法查询可用模型；请显式设置 provider/model"
        )

    loader = available_models_loader or load_available_openai_models
    available_models = [model.strip() for model in loader(base_url, api_key) if str(model).strip()]
    if not available_models:
        raise RuntimeError(
            f"文档页 AI model '{raw_model}' 缺少 provider 前缀，且当前无法查询可用模型；请显式设置 provider/model"
        )

    # 当前 sidecar 走 Responses API，优先选择网关上可用的 responses provider 前缀。
    responses_model = f"ju/{raw_model}"
    if responses_model in available_models:
        return responses_model

    suffix_matches = sorted(
        {
            model
            for model in available_models
            if model.endswith(f"/{raw_model}")
        }
    )
    if len(suffix_matches) == 1:
        return suffix_matches[0]
    if len(suffix_matches) > 1:
        raise RuntimeError(
            f"文档页 AI model '{raw_model}' 存在多个候选，请使用 provider/model 显式指定"
        )

    raise RuntimeError(
        f"未找到可用的文档页 AI model '{raw_model}'；请设置 MNOTE_AI_DEFAULT_MODEL 或在请求中显式传入 provider/model"
    )


def resolve_document_agent_request_model(
    *,
    request_model: str | None,
    request_model_key: str | None,
    configured_default_model: str | None,
    env: Mapping[str, str] | None = None,
    available_models_loader: Callable[[str, str], list[str] | tuple[str, ...]] | None = None,
) -> str:
    direct_model_key = str(request_model_key or "").strip()
    if direct_model_key:
        spec = _document_ai_model_spec_by_key(direct_model_key)
        if spec is not None:
            return spec.gateway_model
        return direct_model_key

    direct_request_model = str(request_model or "").strip()
    if direct_request_model:
        spec = _document_ai_model_spec_by_key(direct_request_model)
        if spec is not None:
            return spec.gateway_model

    return resolve_document_agent_model(
        request_model=request_model,
        configured_default_model=configured_default_model,
        env=env,
        available_models_loader=available_models_loader,
    )


def make_id(prefix: str) -> str:
    return f"{prefix}_{uuid4().hex}"


def to_sse_frame(event: str, data: Any) -> str:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


class DocumentAiMessage(BaseModel):
    role: Literal["user", "assistant"]
    content: str


class DocumentAiContextPayload(BaseModel):
    document_id: str | None = Field(default=None, alias="documentId")
    document_blocks: Any = Field(default=None, alias="documentBlocks")
    node: Any = None
    subtree: Any = None
    outline: Any = None
    evidence: Any = None
    page_options: Any = Field(default=None, alias="pageOptions")
    editor_runtime_page_options: Any = Field(
        default=None,
        alias="editorRuntimePageOptions",
    )

    model_config = {
        "populate_by_name": True,
    }


class DocumentAiRunRequest(BaseModel):
    user_id: str = Field(alias="userId")
    session_id: str | None = Field(default=None, alias="sessionId")
    model: str | None = None
    model_key: str | None = Field(default=None, alias="modelKey")
    profile_id: str | None = Field(default=None, alias="profileId")
    max_steps: int = Field(default=10, alias="maxSteps")
    messages: list[DocumentAiMessage]
    context: DocumentAiContextPayload

    model_config = {
        "populate_by_name": True,
    }


class InsertBlockSpec(BaseModel):
    block_type: Literal["paragraph", "heading"] = Field(
        default="paragraph",
        alias="type",
    )
    text: str = ""
    level: int = 2

    model_config = {
        "populate_by_name": True,
    }


@dataclass
class DocumentAiEmitter:
    queue: asyncio.Queue[tuple[str, Any] | None] = field(
        default_factory=asyncio.Queue,
    )
    tool_count: int = 0
    _tool_sequence: int = 0

    async def emit(self, event: str, data: Any) -> None:
        await self.queue.put((event, data))

    async def close(self) -> None:
        await self.queue.put(None)

    def next_tool_id(self, tool: str) -> str:
        self._tool_sequence += 1
        return f"agents_{tool}_{self._tool_sequence}"


@dataclass
class DocumentAiRunContext:
    user_id: str
    request_id: str
    trace_id: str
    session_id: str | None
    workspace_id: str | None
    document_id: str | None
    document_blocks: Any
    node: Any
    subtree: Any
    outline: Any
    evidence: Any
    page_options: Any
    editor_runtime_page_options: Any
    forward_headers: dict[str, str]
    emitter: DocumentAiEmitter
    bridge: "LocalBridgeRuntime"


@dataclass(frozen=True)
class BridgeRuntimeInvocation:
    command: str
    args: tuple[str, ...]
    cwd: str


@lru_cache(maxsize=1)
def resolve_bridge_repo_root() -> Path:
    repo_root = Path(__file__).resolve().parents[3]
    manifest = repo_root / "rust" / "Cargo.toml"
    if not manifest.exists():
        raise RuntimeError(f"未找到 Rust bridge manifest: {manifest}")
    return repo_root


@lru_cache(maxsize=1)
def resolve_bridge_runtime_invocation() -> BridgeRuntimeInvocation:
    repo_root = resolve_bridge_repo_root()
    explicit_bin = str(os.environ.get("MNOTE_RUST_BRIDGE_BIN", "")).strip()
    if explicit_bin:
        return BridgeRuntimeInvocation(
            command=explicit_bin,
            args=(),
            cwd=str(repo_root),
        )

    for candidate in (
        repo_root / "rust" / "target" / "debug" / "bridge-runtime",
        repo_root / "rust" / "target" / "release" / "bridge-runtime",
    ):
        if candidate.exists():
            return BridgeRuntimeInvocation(
                command=str(candidate),
                args=(),
                cwd=str(repo_root),
            )

    return BridgeRuntimeInvocation(
        command="cargo",
        args=(
            "run",
            "--quiet",
            "--manifest-path",
            str(repo_root / "rust" / "Cargo.toml"),
            "-p",
            "bridge-runtime",
            "--",
        ),
        cwd=str(repo_root),
    )


async def run_bridge_runtime(input_payload: dict[str, Any]) -> dict[str, Any]:
    invocation = resolve_bridge_runtime_invocation()
    payload = json.dumps(input_payload, ensure_ascii=False).encode("utf-8")

    try:
        process = await asyncio.create_subprocess_exec(
            invocation.command,
            *invocation.args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=invocation.cwd,
            env={
                **os.environ,
                "CARGO_TERM_COLOR": "never",
            },
        )
    except Exception as error:
        raise RuntimeError(f"bridge-runtime 启动失败: {error}") from error

    try:
        stdout, stderr = await asyncio.wait_for(
            process.communicate(payload),
            timeout=30.0,
        )
    except asyncio.TimeoutError as error:
        process.kill()
        with contextlib.suppress(ProcessLookupError):
            await process.wait()
        raise RuntimeError("bridge-runtime 执行超时") from error

    stdout_text = stdout.decode("utf-8", errors="replace").strip()
    stderr_text = stderr.decode("utf-8", errors="replace").strip()
    if not stdout_text:
        raise RuntimeError(stderr_text or "bridge-runtime 未返回任何结果")

    try:
        response = json.loads(stdout_text)
    except json.JSONDecodeError as error:
        raise RuntimeError(f"bridge-runtime 返回非法 JSON: {error}") from error

    if not isinstance(response, dict):
        raise RuntimeError("bridge-runtime 返回格式非法")

    if process.returncode != 0 or not bool(response.get("ok")):
        error_payload = response.get("error")
        if isinstance(error_payload, dict):
            message = str(error_payload.get("message", "")).strip()
            if message:
                raise RuntimeError(message)
        raise RuntimeError(
            stderr_text or f"bridge-runtime 执行失败(code={process.returncode})"
        )

    return response


def infer_tool_invocation_kind(tool: str) -> Literal["command", "query", "job"]:
    if tool in {"doc_get", "doc_find"}:
        return "query"
    return "command"


class LocalBridgeRuntime:
    async def execute_tool(
        self,
        *,
        tool: str,
        context: DocumentAiRunContext,
        invocation_kind: Literal["command", "query", "job"] | None = None,
        args_json: dict[str, Any],
        data: Any = None,
        target: dict[str, Any] | None = None,
        reason: str | None = None,
        refs: list[str] | None = None,
    ) -> Any:
        payload = {
            "kind": "tool",
            "context": {
                "deploymentId": None,
                "projectId": None,
                "workspaceId": context.workspace_id,
                "requestId": context.request_id,
                "traceId": context.trace_id,
                "actor": {
                    "actorType": "user",
                    "actorId": context.user_id,
                    "sessionId": context.session_id,
                },
                "source": {
                    "channel": "ai_orchestrator",
                    "client": "wolai-backend",
                },
                "tenantId": None,
                "authToken": None,
                "idempotencyKey": None,
                "validateOnly": False,
                "dryRun": False,
            },
            "tool": {
                "tool": tool,
                "kind": invocation_kind or infer_tool_invocation_kind(tool),
                "mode": "result",
                "argsJson": args_json,
                "target": {
                    "workspaceId": context.workspace_id,
                    "pageId": target.get("pageId") if target else None,
                    "blockId": target.get("blockId") if target else None,
                }
                if target
                else None,
                "reason": reason,
                "refs": refs or [],
            },
        }
        if data is not None:
            payload["data"] = data

        response = await run_bridge_runtime(payload)
        if "result" not in response:
            raise RuntimeError("bridge-runtime 未返回 result")
        return response["result"]


def _safe_json_preview(label: str, value: Any, limit: int) -> str | None:
    if value is None:
        return None
    try:
        text = json.dumps(value, ensure_ascii=False)
    except TypeError:
        return f"{label}=provided"
    return f"{label}={text[:limit]}"


def _extract_inline_text(block: Any) -> str:
    if not isinstance(block, dict):
        return ""
    content = block.get("content")
    if not isinstance(content, list):
        return ""
    parts: list[str] = []
    for node in content:
        if isinstance(node, dict) and isinstance(node.get("text"), str):
            parts.append(node["text"])
    return "".join(parts).strip()


def _normalize_blocks(value: Any) -> list[Any]:
    if isinstance(value, list):
        return value
    if isinstance(value, dict) and isinstance(value.get("blocks"), list):
        return value["blocks"]
    return []


def _summarize_blocks(value: Any, limit: int = 24) -> list[dict[str, Any]]:
    root_blocks = _normalize_blocks(value)
    queue: deque[tuple[Any, int]] = deque((block, 0) for block in root_blocks)
    summaries: list[dict[str, Any]] = []
    while queue and len(summaries) < limit:
        block, depth = queue.popleft()
        if not isinstance(block, dict):
            continue
        block_id = str(block.get("id") or "").strip()
        block_type = str(block.get("type") or "unknown").strip()
        children = block.get("children")
        child_list = children if isinstance(children, list) else []
        if block_id:
            summaries.append(
                {
                    "id": block_id,
                    "type": block_type,
                    "depth": depth,
                    "text": _extract_inline_text(block)[:120],
                    "childCount": len(child_list),
                }
            )
        for child in child_list:
            queue.append((child, depth + 1))
    return summaries


def build_document_agent_instructions(
    request: DocumentAiRunRequest,
) -> str:
    profile = resolve_document_ai_profile(request.profile_id)
    lines = [
        "你是 mnote 的文档页 AI 编写代理。请始终使用简体中文。",
        profile.instructions,
        "你的工作对象只限当前文档页，不要假设你能跨页批量操作，也不要创建新的产品契约。",
        "需要理解当前页内容时，先用 doc_get 或 doc_find；需要写入时，只用 doc_insert_blocks、doc_replace_range、slash_run。",
        "只有在用户明确要求改标题时，才允许用 slash_run，并且只允许对当前 documentId 生成 /rename 命令。",
        "如果工具没有返回成功结果，不要声称页面已经写入完成。",
    ]

    latest_user_message = next(
        (
            message.content.strip()
            for message in reversed(request.messages)
            if message.role == "user" and message.content.strip()
        ),
        "",
    )
    normalized_latest_user_message = latest_user_message.lower()
    hard_rules: list[str] = []
    if request.context.document_id and (
        "标题" in latest_user_message or "rename" in normalized_latest_user_message
    ):
        hard_rules.extend(
            [
                "当前请求包含标题修改意图；你必须调用一次 slash_run，不允许只返回文字说明。",
                f"slash_run 的 text 必须严格使用这个格式：/rename {request.context.document_id} <新标题>。",
                "只有在 slash_run 返回成功后，才能用一句简短中文确认标题已修改。",
            ]
        )
    if "blockid" in normalized_latest_user_message and (
        "改写" in latest_user_message or "替换" in latest_user_message
    ):
        hard_rules.extend(
            [
                "当前请求要求按 blockId 改写正文；你必须直接调用 doc_replace_range，不要只返回文字说明。",
                "如果用户消息里已经给了 blockId 和目标文本，就不要先调用 doc_find。",
                "只有在 doc_replace_range 返回成功后，才能用一句简短中文确认正文已修改。",
            ]
        )
    if ("插入" in latest_user_message or "新增" in latest_user_message) and (
        "正文" in latest_user_message or "段" in latest_user_message
    ):
        hard_rules.extend(
            [
                "当前请求要求新增正文；你必须调用 doc_insert_blocks，不允许只返回文字说明。",
                "doc_insert_blocks 至少要插入一个 paragraph 或 heading block；除非用户明确要求，不要顺手改标题。",
                "只有在 doc_insert_blocks 返回成功后，才能用一句简短中文确认正文已写入。",
            ]
        )
    if hard_rules:
        lines.append("当前请求的强制执行规则：\n- " + "\n- ".join(hard_rules))

    context = request.context
    lines.append(f"profileId={profile.id}")
    if context.document_id:
        lines.append(f"documentId={context.document_id}")
    lines.append(
        "documentBlockSummaries="
        + json.dumps(
            _summarize_blocks(context.document_blocks),
            ensure_ascii=False,
        )
    )
    for label, value, limit in (
        ("pageOptions", context.page_options, 1600),
        ("editorRuntimePageOptions", context.editor_runtime_page_options, 1200),
        ("node", context.node, 1600),
        ("subtree", context.subtree, 2400),
        ("outline", context.outline, 2400),
        ("evidence", context.evidence, 2400),
    ):
        preview = _safe_json_preview(label, value, limit)
        if preview:
            lines.append(preview)
    return "\n\n".join(lines).strip()


def _trim_messages(messages: list[DocumentAiMessage]) -> list[dict[str, str]]:
    return [
        {"role": message.role, "content": message.content}
        for message in messages[-24:]
        if message.content.strip()
    ]


async def _emit_tool_event(
    ctx: DocumentAiRunContext,
    tool: str,
    args_json: dict[str, Any],
    execute: Callable[[], Any],
) -> Any:
    tool_id = ctx.emitter.next_tool_id(tool)
    await ctx.emitter.emit(
        "tool_call",
        {
            "id": tool_id,
            "tool": tool,
            "args": args_json,
        },
    )

    started = time.perf_counter()
    try:
        result = await execute()
    except Exception as error:
        elapsed_ms = max(0, round((time.perf_counter() - started) * 1000))
        await ctx.emitter.emit(
            "tool_result",
            {
                "id": tool_id,
                "tool": tool,
                "ok": False,
                "ms": elapsed_ms,
                "result": {"error": str(error)},
            },
        )
        raise

    ctx.emitter.tool_count += 1
    elapsed_ms = max(0, round((time.perf_counter() - started) * 1000))
    await ctx.emitter.emit(
        "tool_result",
        {
            "id": tool_id,
            "tool": tool,
            "ok": True,
            "ms": elapsed_ms,
            "result": result,
        },
    )
    return result


def build_document_tools():
    if not sdk_available():
        raise RuntimeError("openai-agents 未安装，无法构建文档页 tools")

    @function_tool
    async def doc_get(
        wrapper: RunContextWrapper[DocumentAiRunContext],
        maxBlocks: int = 80,
    ) -> Any:
        """读取当前页的块摘要。需要先理解页面结构时调用。"""

        ctx = wrapper.context
        args_json = {"maxBlocks": max(10, min(int(maxBlocks), 240))}
        return await _emit_tool_event(
            ctx,
            "doc_get",
            args_json,
            lambda: ctx.bridge.execute_tool(
                tool="doc_get",
                context=ctx,
                invocation_kind="query",
                args_json=args_json,
                data=ctx.document_blocks,
                target={"pageId": ctx.document_id},
            ),
        )

    @function_tool
    async def doc_find(
        wrapper: RunContextWrapper[DocumentAiRunContext],
        query: str,
        maxResults: int = 8,
    ) -> Any:
        """按文本查找当前页的块。准备改写某段内容前先调用。"""

        ctx = wrapper.context
        query_text = str(query).strip()
        if not query_text:
            raise RuntimeError("doc_find 缺少 query")
        args_json = {
            "query": query_text,
            "maxResults": max(1, min(int(maxResults), 30)),
        }
        return await _emit_tool_event(
            ctx,
            "doc_find",
            args_json,
            lambda: ctx.bridge.execute_tool(
                tool="doc_find",
                context=ctx,
                invocation_kind="query",
                args_json=args_json,
                data=ctx.document_blocks,
                target={"pageId": ctx.document_id},
            ),
        )

    @function_tool
    async def doc_insert_blocks(
        wrapper: RunContextWrapper[DocumentAiRunContext],
        blocks: list[InsertBlockSpec],
        afterBlockId: str | None = None,
        beforeBlockId: str | None = None,
    ) -> Any:
        """在当前页插入新块。新增标题或段落时调用。"""

        ctx = wrapper.context
        if not blocks:
            raise RuntimeError("doc_insert_blocks 缺少 blocks")
        args_json = {
            "blocks": [
                {
                    "type": block.block_type,
                    "text": block.text,
                    "level": block.level,
                }
                for block in blocks[:20]
            ],
        }
        if afterBlockId:
            args_json["afterBlockId"] = afterBlockId
        if beforeBlockId:
            args_json["beforeBlockId"] = beforeBlockId
        return await _emit_tool_event(
            ctx,
            "doc_insert_blocks",
            args_json,
            lambda: ctx.bridge.execute_tool(
                tool="doc_insert_blocks",
                context=ctx,
                invocation_kind="command",
                args_json=args_json,
                data=ctx.document_blocks,
                target={"pageId": ctx.document_id},
            ),
        )

    @function_tool
    async def doc_replace_range(
        wrapper: RunContextWrapper[DocumentAiRunContext],
        blockId: str,
        text: str,
        mode: Literal["replace", "append", "prepend"] = "replace",
    ) -> Any:
        """改写当前页某个块的文本。"""

        ctx = wrapper.context
        block_id = str(blockId).strip()
        content = str(text).strip()
        if not block_id or not content:
            raise RuntimeError("doc_replace_range 缺少 blockId 或 text")
        safe_mode = mode if mode in {"replace", "append", "prepend"} else "replace"
        args_json = {
            "blockId": block_id,
            "text": content,
            "mode": safe_mode,
        }
        return await _emit_tool_event(
            ctx,
            "doc_replace_range",
            args_json,
            lambda: ctx.bridge.execute_tool(
                tool="doc_replace_range",
                context=ctx,
                invocation_kind="command",
                args_json=args_json,
                data=ctx.document_blocks,
                target={"pageId": ctx.document_id, "blockId": block_id},
            ),
        )

    @function_tool
    async def slash_run(
        wrapper: RunContextWrapper[DocumentAiRunContext],
        text: str,
    ) -> Any:
        """执行当前页允许的斜杠命令。当前只允许改当前页标题。"""

        ctx = wrapper.context
        command_text = str(text).strip()
        if not command_text.startswith("/rename "):
            raise RuntimeError("当前只允许 /rename 当前页标题")
        if ctx.document_id and f"/rename {ctx.document_id} " not in command_text:
            raise RuntimeError("slash_run 只允许改当前文档页标题")
        args_json = {"text": command_text}
        return await _emit_tool_event(
            ctx,
            "slash_run",
            args_json,
            lambda: ctx.bridge.execute_tool(
                tool="slash_run",
                context=ctx,
                invocation_kind="command",
                args_json=args_json,
                data={"source": "ai-orchestrator"},
                target={"pageId": ctx.document_id},
            ),
        )

    return [
        doc_get,
        doc_find,
        doc_insert_blocks,
        doc_replace_range,
        slash_run,
    ]


def build_document_agent(
    request: DocumentAiRunRequest,
):
    if not sdk_available():
        raise RuntimeError("openai-agents 未安装，无法构建文档页 agent")

    kwargs: dict[str, Any] = {
        "name": "mnote-document-page-writer",
        "instructions": build_document_agent_instructions(request),
        "tools": build_document_tools(),
    }
    kwargs["model"] = resolve_document_agent_request_model(
        request_model=request.model,
        request_model_key=request.model_key,
        configured_default_model=settings.mnote_ai_default_model,
    )
    return Agent[DocumentAiRunContext](**kwargs)


def build_forward_headers(source_headers: dict[str, str]) -> dict[str, str]:
    allowed = {
        "authorization",
        "cookie",
        "x-request-id",
        "x-trace-id",
        "x-session-id",
        "x-mnote-workspace-id",
        "x-mnote-source-channel",
        "x-mnote-source-client",
        "x-mnote-actor-id",
        "x-mnote-actor-type",
        "user-agent",
    }
    return {
        key: value
        for key, value in source_headers.items()
        if key.lower() in allowed and value
    }


def build_session(session_id: str | None):
    if not sdk_available() or not session_id:
        return None

    raw_path = settings.openai_agents_session_db_path.strip()
    db_path = (
        Path(raw_path)
        if raw_path
        else Path(__file__).resolve().parents[2] / ".data" / "openai-agents-sessions.sqlite3"
    )
    db_path.parent.mkdir(parents=True, exist_ok=True)
    return SQLiteSession(session_id, str(db_path))


async def run_document_agent_stream(
    request: DocumentAiRunRequest,
    *,
    source_headers: dict[str, str],
) -> AsyncGenerator[str, None]:
    if not sdk_available():
        raise RuntimeError("openai-agents 未安装")
    if not has_configured_openai_api_key():
        raise RuntimeError("后端未配置 OPENAI_API_KEY")
    effective_api_key = get_effective_openai_api_key()
    if effective_api_key:
        os.environ.setdefault("OPENAI_API_KEY", effective_api_key)

    emitter = DocumentAiEmitter()
    run_context = DocumentAiRunContext(
        user_id=request.user_id,
        request_id=make_id("req"),
        trace_id=make_id("trace"),
        session_id=request.session_id,
        workspace_id=source_headers.get("x-mnote-workspace-id") or None,
        document_id=request.context.document_id,
        document_blocks=request.context.document_blocks,
        node=request.context.node,
        subtree=request.context.subtree,
        outline=request.context.outline,
        evidence=request.context.evidence,
        page_options=request.context.page_options,
        editor_runtime_page_options=request.context.editor_runtime_page_options,
        forward_headers=build_forward_headers(source_headers),
        emitter=emitter,
        bridge=LocalBridgeRuntime(),
    )

    async def worker() -> None:
        try:
            agent = build_document_agent(request)
            run_config = RunConfig(
                model_provider=build_document_model_provider(
                    request_model=request.model,
                    configured_default_model=settings.mnote_ai_default_model,
                )
            )
            session = build_session(request.session_id)
            run_input = _trim_messages(request.messages)
            if not run_input:
                raise RuntimeError("缺少 messages")

            max_turns = max(1, min(int(request.max_steps), 12))
            with trace("mnote-document-ai-run", group_id=run_context.trace_id):
                result = await Runner.run(
                    agent,
                    input=run_input,
                    context=run_context,
                    session=session,
                    max_turns=max_turns,
                    run_config=run_config,
                )

            text = str(result.final_output or "").strip() or "（无输出）"
            await emitter.emit("assistant_message", {"text": text})
            await emitter.emit(
                "completion",
                {
                    "ok": True,
                    "text": text,
                    "steps": max(1, emitter.tool_count or 1),
                },
            )
        except Exception as error:
            await emitter.emit(
                "error",
                {
                    "ok": False,
                    "message": str(error),
                },
            )
        finally:
            await emitter.close()

    task = asyncio.create_task(worker())
    try:
        while True:
            item = await emitter.queue.get()
            if item is None:
                break
            event, data = item
            yield to_sse_frame(event, data)
    finally:
        if not task.done():
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
