"""LLMWiki retrieval helpers for normal DeerFlow chat runs.

The public chat API accepts DeerFlow-local LLMWiki knowledge-base mapping ids.
When WeKnora mode is enabled, this module resolves those local ids under the
current DeerFlow user's permissions, searches WeKnora once before the model
call, and attaches both:

* a hidden prompt with numbered context snippets, and
* structured citation sources for the existing reference panel.

The helper lives in ``app.gateway`` because it needs Gateway auth state and the
DeerFlow-owned mapping store.  Harness code remains app-import free.
"""

from __future__ import annotations

import logging
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any
from urllib.parse import urlencode

from fastapi import HTTPException, Request
from langchain_core.messages import HumanMessage

from app.gateway.deps import get_optional_user_from_request
from app.gateway.utils import sanitize_log_param
from deerflow.config.agents_config import load_agent_config, validate_agent_id
from deerflow.integrations.weknora.client import WeKnoraError
from deerflow.integrations.weknora.runtime import build_weknora_client, get_resolved_llmwiki_runtime
from deerflow.persistence.llmwiki import LlmWikiStore
from deerflow.runtime.user_context import get_effective_user_id

logger = logging.getLogger(__name__)

_MAX_SOURCES = 8
_MAX_SOURCE_CHARS = 1400


@dataclass(frozen=True)
class LlmWikiRagContext:
    knowledge_base_ids: list[str]
    source_count: int
    prompt: str
    sources: list[dict[str, Any]]
    query: str


def _store(request: Request) -> LlmWikiStore | None:
    return getattr(request.app.state, "llmwiki_store", None)


async def _actor(request: Request) -> tuple[str, bool]:
    user = await get_optional_user_from_request(request)
    if user is None:
        return get_effective_user_id(), True
    return str(user.id), getattr(user, "system_role", None) == "admin"


def _as_id_list(value: Any) -> list[str]:
    if value is None:
        return []
    if isinstance(value, str):
        parts = [value]
    elif isinstance(value, (list, tuple, set)):
        parts = list(value)
    else:
        return []
    return list(dict.fromkeys(str(item).strip() for item in parts if str(item).strip()))


def _body_selected_ids(body: Any) -> list[str]:
    for source in (getattr(body, "context", None), getattr(body, "metadata", None)):
        if isinstance(source, Mapping) and "llmwiki_knowledge_base_ids" in source:
            return _as_id_list(source.get("llmwiki_knowledge_base_ids"))
    return []


def _assistant_agent_id(body: Any) -> str | None:
    assistant_id = validate_agent_id(getattr(body, "assistant_id", None))
    if assistant_id:
        return assistant_id
    context = getattr(body, "context", None)
    if isinstance(context, Mapping):
        return validate_agent_id(str(context.get("agent_id") or "")) if context.get("agent_id") else None
    metadata = getattr(body, "metadata", None)
    if isinstance(metadata, Mapping):
        return validate_agent_id(str(metadata.get("agent_id") or "")) if metadata.get("agent_id") else None
    return None


def _agent_bound_ids(body: Any) -> list[str]:
    agent_id = _assistant_agent_id(body)
    if not agent_id:
        return []
    try:
        config = load_agent_config(agent_id)
    except (FileNotFoundError, ValueError):
        return []
    return _as_id_list(config.llmwiki_knowledge_base_ids if config else None)


def resolve_requested_llmwiki_knowledge_base_ids(body: Any) -> list[str]:
    """Return selected ids plus custom-agent defaults, de-duplicated in order."""

    return list(dict.fromkeys([*_agent_bound_ids(body), *_body_selected_ids(body)]))


def llmwiki_selection_requested(body: Any) -> bool:
    return bool(resolve_requested_llmwiki_knowledge_base_ids(body))


def _latest_human_query(graph_input: dict[str, Any]) -> str:
    messages = graph_input.get("messages")
    if not isinstance(messages, list):
        return ""
    for message in reversed(messages):
        if isinstance(message, HumanMessage):
            return _message_text(getattr(message, "content", ""))
        if isinstance(message, dict) and str(message.get("type") or message.get("role") or "") in {"human", "user"}:
            return _message_text(message.get("content"))
    return ""


def _message_text(content: Any) -> str:
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, Mapping):
                text = item.get("text") or item.get("content")
                if isinstance(text, str):
                    parts.append(text)
        return "\n".join(part.strip() for part in parts if part and part.strip()).strip()
    return str(content or "").strip()


def _truncate(text: str, max_chars: int) -> str:
    value = text.strip()
    if len(value) <= max_chars:
        return value
    return f"{value[: max_chars - 1].rstrip()}…"


def _source_from_result(result: dict[str, Any], mapping: dict[str, Any], index: int) -> dict[str, Any]:
    chunk_id = str(result.get("chunk_id") or "")
    title = str(result.get("title") or result.get("filename") or mapping.get("name") or f"LLMWiki 片段 {index}")
    content = str(result.get("content") or "")
    params = urlencode({"knowledge_base_id": mapping["id"], "chunk_id": chunk_id})
    raw = {
        **result,
        "knowledge_base_id": mapping["id"],
        "knowledge_base_name": mapping.get("name") or "",
        "source": mapping.get("name") or "LLMWiki",
        "url": f"/weknora-source-preview?{params}",
    }
    return {
        "index": index,
        "id": chunk_id or str(result.get("knowledge_id") or f"{mapping['id']}:{index}"),
        "title": title,
        "type": "llmwiki",
        "snippet": _truncate(content, 360),
        "content": content,
        "source": mapping.get("name") or "LLMWiki",
        "score": float(result.get("score") or 0),
        "url": f"/weknora-source-preview?{params}",
        "skillName": "llmwiki_search",
        "mode": "citation",
        "raw": raw,
    }


def _build_prompt(query: str, sources: list[dict[str, Any]]) -> str:
    if not sources:
        return (
            "当前用户选择了 LLMWiki 知识库，但 DeerFlow 未检索到可用片段。"
            "请直接说明没有检索到可用背景资料，不要编造知识库内容或引用编号。"
        )

    lines = [
        "当前用户选择了 LLMWiki 知识库。DeerFlow 已在模型回答前自动检索到以下背景资料。",
        "请优先依据这些资料回答；如果资料不足，请明确说明不足之处，不要编造。",
        "引用规则：正文关键结论必须使用 [1]、[2] 这样的编号引用下方资料；只能引用存在的编号；不要在正文末尾重复输出完整参考文献列表，前端会展示引用卡片。",
        "",
        f"用户问题：{query}",
        "",
        f"<llmwiki_reference_batch source_count=\"{len(sources)}\">",
    ]
    for source in sources:
        meta = " / ".join(
            str(item)
            for item in (
                source.get("source"),
                f"score={source.get('score'):.4f}" if isinstance(source.get("score"), float) else None,
            )
            if item
        )
        lines.append(f"[{source['index']}] {source.get('title') or 'LLMWiki 片段'}")
        if meta:
            lines.append(f"来源：{meta}")
        lines.append(_truncate(str(source.get("content") or source.get("snippet") or ""), _MAX_SOURCE_CHARS))
        lines.append("")
    lines.append("</llmwiki_reference_batch>")
    return "\n".join(lines).strip()


async def build_llmwiki_rag_context(
    *,
    request: Request,
    body: Any,
    graph_input: dict[str, Any],
) -> LlmWikiRagContext:
    selected_ids = resolve_requested_llmwiki_knowledge_base_ids(body)
    if not selected_ids:
        return LlmWikiRagContext([], 0, "", [], "")

    runtime = get_resolved_llmwiki_runtime(getattr(request.app.state, "config", None))
    if not runtime.weknora_enabled:
        return LlmWikiRagContext(selected_ids, 0, "", [], "")
    store = _store(request)
    if store is None:
        return LlmWikiRagContext(selected_ids, 0, "", [], "")

    query = _latest_human_query(graph_input)
    if not query:
        return LlmWikiRagContext(selected_ids, 0, "", [], "")

    user_id, is_admin = await _actor(request)
    rows: list[dict[str, Any]] = []
    for mapping_id in selected_ids:
        row = await store.get_authorized(mapping_id, user_id, write=False, is_admin=is_admin)
        if row is None:
            raise HTTPException(status_code=404, detail=f"Knowledge base {mapping_id} not found")
        rows.append(row)
    remote_to_local = {str(row["weknora_id"]): row for row in rows}

    try:
        results = await build_weknora_client(runtime).search(query, [str(row["weknora_id"]) for row in rows])
    except WeKnoraError as exc:
        logger.warning(
            "LLMWiki pre-retrieval failed for user=%s kb_count=%s: %s",
            sanitize_log_param(user_id),
            len(rows),
            exc,
        )
        return LlmWikiRagContext(
            selected_ids,
            0,
            f"当前用户选择了 LLMWiki 知识库，但自动检索失败：{exc}。请说明暂时无法检索知识库背景资料，不要编造引用。",
            [],
            query,
        )

    sources: list[dict[str, Any]] = []
    for result in results:
        mapping = remote_to_local.get(str(result.get("knowledge_base_id") or ""))
        if mapping is None or not str(result.get("content") or "").strip():
            continue
        sources.append(_source_from_result(result, mapping, len(sources) + 1))
        if len(sources) >= _MAX_SOURCES:
            break
    return LlmWikiRagContext(selected_ids, len(sources), _build_prompt(query, sources), sources, query)


def attach_llmwiki_rag_to_graph_input(
    graph_input: dict[str, Any],
    rag: LlmWikiRagContext,
) -> dict[str, Any]:
    if not rag.prompt and not rag.sources:
        return graph_input
    messages = graph_input.get("messages")
    if not isinstance(messages, list) or not messages:
        return graph_input
    new_messages = list(messages)
    for index, message in enumerate(new_messages):
        if not isinstance(message, HumanMessage):
            continue
        extra = dict(getattr(message, "additional_kwargs", None) or {})
        existing = str(extra.get("prompt_prefix") or "").strip()
        extra["prompt_prefix"] = f"{existing}\n\n{rag.prompt}".strip() if existing else rag.prompt
        extra["llmwiki_knowledge_base_ids"] = rag.knowledge_base_ids
        extra["llmwiki_reference_query"] = rag.query
        if rag.sources:
            extra["llmwiki_reference_sources"] = rag.sources
        new_messages[index] = message.model_copy(update={"additional_kwargs": extra})
        break
    return {**graph_input, "messages": new_messages}