"""Knowledge base query workflow — search and review knowledge.

Reuses:
    - ``src.graphrag.unified.rag_search(action, query, ...)``
    - ``src.graphrag.unified.rag_index(action, ...)``

IKE integration (optional):
    - P1 Knowledge Base: Hybrid search (vector + graph + lexical) via IKE
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from databridge_core.sdk.core.context import SDKContext
from databridge_core.sdk.core.exceptions import NotImplementedStubError
from databridge_core.sdk.core.types import CompiledArtifact, Result
from databridge_core.sdk.workflows.base import BaseWorkflow
from databridge_core.sdk.workflows.result import WorkflowResult
from databridge_core.sdk.workflows.step import workflow_step


class KBQueryWorkflow(BaseWorkflow):
    """Search and query the knowledge base via GraphRAG.

    Supports forensic, hierarchy, catalog, and semantic search modes.

    When IKE is available, also queries IKE's P1 hybrid knowledge base
    (vector + graph + lexical with RRF fusion) and merges results.

    Reuses:
        - ``src.graphrag.unified.rag_search(action, query, ...)``
        - ``src.graphrag.unified.rag_index(action, ...)``
        - ``src.graphrag.unified.rag_ops(action, ...)``
    """

    name = "knowledge_base.kb_query"
    category = "knowledge_base"
    description = "Search and query the knowledge base via GraphRAG"
    required_params = ["query"]
    optional_params = {
        "search_type": "forensic",
        "max_results": 10,
        "include_lineage": True,
        "ike_hybrid_search": True,
    }

    def validate(self, params: Dict[str, Any]) -> Dict[str, Any]:
        self._check_required(params)
        return self._apply_defaults(params)

    def execute(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> Result:
        steps = [self._search(params, ctx)]
        # IKE P1: Augment with hybrid KB search if available
        if params.get("ike_hybrid_search") and ctx and ctx.ike.available:
            steps.append(self._ike_hybrid_search(params, ctx))
        return WorkflowResult.from_steps(self.name, steps)

    def compile_to_snowflake(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> List[CompiledArtifact]:
        raise NotImplementedStubError(
            self.name,
            "compile_to_snowflake",
            "Knowledge base queries are runtime operations; "
            "for Snowflake deployment, use Cortex Search Service",
        )

    @workflow_step("search", "Search knowledge base")
    def _search(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # REUSE: src.graphrag.unified.rag_search(
        #     action=params["search_type"],
        #     query=params["query"],
        #     max_results=params["max_results"],
        #     include_lineage=params["include_lineage"],
        # )
        raise NotImplementedStubError(
            self.name,
            "search",
            "src.graphrag.unified.rag_search(action, query, max_results, include_lineage)",
        )

    @workflow_step("ike_hybrid_search", "Search IKE hybrid KB (vector+graph+lexical)")
    def _ike_hybrid_search(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # IKE P1: Hybrid search combining vector similarity, graph
        # traversal, and lexical matching with RRF rank fusion.
        # Results augment the GraphRAG search from the prior step.
        #
        # ctx.ike.query_knowledge(
        #     query=params["query"],
        #     top_k=params["max_results"],
        # )
        raise NotImplementedStubError(
            self.name,
            "ike_hybrid_search",
            "ctx.ike.query_knowledge(query, top_k) — IKE P1 hybrid retrieval with RRF",
        )
