# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

"""Knowledge base tools factory module.

Responsible for creating knowledge base search tools and enhancing system prompts.
"""

import logging
from typing import Any, List, Optional

from sqlalchemy.ext.asyncio import AsyncSession

from shared.models.knowledge import (
    KnowledgeBaseToolAccessMode,
    KnowledgeBaseToolsResult,
)
from shared.telemetry.decorators import add_span_event, set_span_attribute, trace_async

logger = logging.getLogger(__name__)


async def prepare_knowledge_base_tools(
    knowledge_base_ids: Optional[list[int]],
    user_id: int,
    db: AsyncSession,
    base_system_prompt: str,
    task_id: Optional[int] = None,
    user_subtask_id: Optional[int] = None,
    is_user_selected: bool = True,
    document_ids: Optional[list[int]] = None,
    model_id: Optional[str] = None,
    context_window: Optional[int] = None,
    model_config: Optional[dict[str, Any]] = None,
    skip_prompt_enhancement: bool = False,
    user_name: Optional[str] = None,
    auth_token: str = "",
    kb_tool_access_mode: str = KnowledgeBaseToolAccessMode.FULL,
) -> KnowledgeBaseToolsResult:
    """Prepare knowledge base tools and enhanced system prompt.

    This function encapsulates the logic for creating knowledge base tools and
    enhancing the system prompt with knowledge base instructions.

    NOTE:
    - Knowledge base metadata prompt (kb_meta_prompt) is generated by Backend and transported
      via the execution request (e.g. `ExecutionRequest.kb_meta_prompt`). chat_shell must not
      build it locally to avoid backend dependency in HTTP mode.

    Returns:
        KnowledgeBaseToolsResult(extra_tools, enhanced_system_prompt, kb_meta_prompt)

        kb_meta_prompt is always an empty string here and should be injected by the caller
        via `dynamic_context` if needed.
    """
    extra_tools: list = []
    enhanced_system_prompt = base_system_prompt

    if not knowledge_base_ids:
        return KnowledgeBaseToolsResult(
            extra_tools=extra_tools,
            enhanced_system_prompt=enhanced_system_prompt,
            kb_meta_prompt="",
        )

    logger.info(
        "[knowledge_factory] Creating KB tools for %d knowledge bases: %s, "
        "is_user_selected=%s, document_ids=%s, model_id=%s, context_window=%s",
        len(knowledge_base_ids),
        knowledge_base_ids,
        is_user_selected,
        document_ids,
        model_id,
        context_window,
    )

    # Import knowledge base tools
    from chat_shell.tools.builtin import (
        KbHeadTool,
        KbLsTool,
        KBToolCallCounter,
        KnowledgeBaseTool,
    )

    restricted_search_only = (
        kb_tool_access_mode == KnowledgeBaseToolAccessMode.RESTRICTED_SEARCH_ONLY
    )

    # Create KnowledgeBaseTool with the specified knowledge bases
    # KB configs (max_calls, exempt_calls, name) are fetched from Backend API
    # Pass user_subtask_id for persisting RAG results to context database
    # Pass document_ids for filtering to specific documents
    # Pass context_window from Model CRD for injection strategy decisions
    kb_tool = KnowledgeBaseTool(
        knowledge_base_ids=knowledge_base_ids,
        document_ids=document_ids or [],
        user_id=user_id,
        user_name=user_name,
        auth_token=auth_token,
        db_session=db,
        user_subtask_id=user_subtask_id,
        model_id=model_id or KnowledgeBaseTool.model_id,
        current_model_name=(model_config or {}).get("model_name"),
        current_model_namespace=(model_config or {}).get("model_namespace")
        or "default",
        context_window=context_window,
        injection_mode="hybrid",
        tool_access_mode=kb_tool_access_mode,
    )
    extra_tools.append(kb_tool)

    if restricted_search_only:
        logger.info(
            "[knowledge_factory] Created restricted KB tool set: knowledge_base_search only"
        )
    else:
        # Create shared call counter for exploration tools (kb_ls and kb_head).
        # knowledge_base_search enforces its own limit internally; we fetch the
        # same KB-configured max_calls here only because kb_ls/kb_head need a
        # shared counter when they are actually exposed.
        try:
            max_calls, _ = kb_tool._get_kb_limits()
        except Exception:
            # KnowledgeBaseTool may be mocked in unit tests; fall back to defaults.
            max_calls = 10
        exploration_call_counter = KBToolCallCounter(max_calls=max_calls)

        # Create exploration tools (kb_ls and kb_head)
        # These are secondary tools for when RAG search doesn't find relevant results
        # They share a call counter to enforce combined call limits
        kb_ls_tool = KbLsTool(
            knowledge_base_ids=knowledge_base_ids,
            db_session=db,
            auth_token=auth_token,
        )
        kb_ls_tool._call_counter = exploration_call_counter

        kb_head_tool = KbHeadTool(
            knowledge_base_ids=knowledge_base_ids,
            user_id=user_id,
            db_session=db,
            user_subtask_id=user_subtask_id,
            auth_token=auth_token,
        )
        kb_head_tool._call_counter = exploration_call_counter

        extra_tools.extend([kb_ls_tool, kb_head_tool])

        logger.info(
            "[knowledge_factory] Created 3 KB tools: knowledge_base_search, kb_ls, kb_head"
        )

    # Skip prompt enhancement if Backend has already added KB prompts (HTTP mode)
    if skip_prompt_enhancement:
        logger.info(
            "[knowledge_factory] Skipping KB prompt enhancement (already done by Backend)"
        )
        return KnowledgeBaseToolsResult(
            extra_tools=extra_tools,
            enhanced_system_prompt=enhanced_system_prompt,
            kb_meta_prompt="",
            kb_tool_access_mode=kb_tool_access_mode,
        )

    # Import shared prompt constants from chat_shell prompts module
    from chat_shell.prompts import (
        KB_PROMPT_NO_RAG,
        KB_PROMPT_RELAXED,
        KB_PROMPT_RESTRICTED_ANALYST,
        KB_PROMPT_STRICT,
    )

    if restricted_search_only:
        kb_instruction = KB_PROMPT_RESTRICTED_ANALYST
        logger.info(
            "[knowledge_factory] Using RESTRICTED_ANALYST mode prompt "
            "(search-only KB access)"
        )
        enhanced_system_prompt = f"{base_system_prompt}{kb_instruction}"
        return KnowledgeBaseToolsResult(
            extra_tools=extra_tools,
            enhanced_system_prompt=enhanced_system_prompt,
            kb_meta_prompt="",
            kb_tool_access_mode=kb_tool_access_mode,
        )

    # Check if any KB has RAG enabled by querying KB info
    has_rag_enabled = await _check_any_kb_has_rag_enabled(
        knowledge_base_ids, auth_token
    )

    # Choose prompt based on RAG availability and user selection mode
    if not has_rag_enabled:
        # No-RAG mode: Use exploration tools only
        kb_instruction = KB_PROMPT_NO_RAG
        logger.info(
            "[knowledge_factory] Using NO_RAG mode prompt (no retriever configured)"
        )
    elif is_user_selected:
        # Strict mode: User explicitly selected KB for this message
        kb_instruction = KB_PROMPT_STRICT
        logger.info(
            "[knowledge_factory] Using STRICT mode prompt (user explicitly selected KB)"
        )
    else:
        # Relaxed mode: KB inherited from task, AI can use general knowledge as fallback
        kb_instruction = KB_PROMPT_RELAXED
        logger.info(
            "[knowledge_factory] Using RELAXED mode prompt (KB inherited from task)"
        )

    # Keep KB prompt template fully static (no kb_meta_list placeholder).
    enhanced_system_prompt = f"{base_system_prompt}{kb_instruction}"

    # kb_meta_prompt is generated by Backend and transported via execution request.
    return KnowledgeBaseToolsResult(
        extra_tools=extra_tools,
        enhanced_system_prompt=enhanced_system_prompt,
        kb_meta_prompt="",
        kb_tool_access_mode=kb_tool_access_mode,
    )


@trace_async(
    span_name="check_any_kb_has_rag_enabled",
    tracer_name="chat_shell.tools.knowledge_factory",
)
async def _check_any_kb_has_rag_enabled(
    knowledge_base_ids: list[int], auth_token: str = ""
) -> bool:
    """
    Check if any of the given knowledge bases have RAG enabled.

    This function queries the Backend API to get KB info and checks if any KB
    has a retriever configured (rag_enabled=True).

    Args:
        knowledge_base_ids: List of knowledge base IDs to check

    Returns:
        True if at least one KB has RAG enabled, False otherwise
    """
    set_span_attribute("knowledge_base_ids.count", len(knowledge_base_ids))

    if not knowledge_base_ids:
        add_span_event("no_knowledge_bases")
        return False

    from chat_shell.core.config import settings

    try:
        import httpx

        add_span_event("querying_backend_api")
        headers = {}
        # Priority: INTERNAL_SERVICE_TOKEN (service-to-service) > auth_token (task token)
        token = getattr(settings, "INTERNAL_SERVICE_TOKEN", "") or auth_token
        if token:
            headers["Authorization"] = f"Bearer {token}"
            set_span_attribute("auth_token_provided", True)
        else:
            set_span_attribute("auth_token_provided", False)

        # Query KB info from Backend API
        url = f"{settings.BACKEND_URL}/api/internal/rag/kb-size"
        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.post(
                url,
                json={"knowledge_base_ids": knowledge_base_ids},
                headers=headers or None,
            )
            response.raise_for_status()
            data = response.json()

            # Check if any KB has rag_enabled=True
            items = data.get("items", [])
            for item in items:
                if item.get("rag_enabled", False):
                    logger.debug(
                        f"[knowledge_factory] KB {item.get('id')} has RAG enabled"
                    )
                    add_span_event("rag_enabled_found")
                    set_span_attribute("rag_enabled", True)
                    set_span_attribute("kb_id_with_rag", item.get("id"))
                    return True

            logger.info(
                f"[knowledge_factory] No KB has RAG enabled among {knowledge_base_ids}"
            )
            add_span_event("no_rag_enabled")
            set_span_attribute("rag_enabled", False)
            return False

    except Exception as e:
        logger.warning(
            f"[knowledge_factory] Failed to check RAG status for KBs {knowledge_base_ids}: {e}. "
            "Assuming RAG is enabled (fallback to normal behavior)."
        )
        add_span_event("check_failed")
        set_span_attribute("error", str(e))
        set_span_attribute("rag_enabled", True)  # fallback
        # On error, assume RAG is enabled to avoid breaking existing functionality
        return True
