import json
import logging

from langchain_core.tools import tool
from pydantic import BaseModel, Field

from agent.prompt_utils import read_tool_prompt
from common.embedding_service import create_embedding_async
from common.starrocks_connection import get_db_connection

logger = logging.getLogger(__name__)

TOP_K_CHUNKS = 5

class KnowledgeSearchInput(BaseModel):
    query: str = Field(
        description="Câu hỏi hoặc nhu cầu tìm kiếm thông tin phi sản phẩm của khách hàng (ví dụ: hỏi chính sách, tra bảng size...). KHÔNG DÙNG ĐỂ TÌM SẢN PHẨM. KHÔNG DÙNG ĐỂ TÌM CỬA HÀNG (dùng canifa_store_search)."
    )


@tool("canifa_knowledge_search", args_schema=KnowledgeSearchInput)
async def canifa_knowledge_search(query: str) -> str:
    """
    (Placeholder docstring - Actual prompt is loaded from file)
    Tra cứu thông tin thương hiệu Canifa (chính sách, KHTT, bảng size...).
    KHÔNG dùng để tìm cửa hàng - dùng canifa_store_search.
    """
    logger.info(f"[Semantic Search] Brand Knowledge query: {query}")
    try:
        # 1. Tạo embedding cho câu hỏi (cached 24h trong Redis)
        query_vector = await create_embedding_async(query)
        if not query_vector:
            return "Xin lỗi, tôi gặp sự cố khi xử lý thông tin. Vui lòng thử lại sau."

        v_str = json.dumps(query_vector)

        # 2. SINGLE QUERY: Tìm top K chunks → lấy full document theo title (1 round-trip)
        sql = f"""
            WITH top_chunks AS (
                SELECT json_query(metadata, '$.title') AS title
                FROM shared_source.chatbot_rsa_knowledge
                ORDER BY approx_cosine_similarity(embedding, {v_str}) DESC
                LIMIT {TOP_K_CHUNKS}
            )
            SELECT k.content, k.metadata
            FROM shared_source.chatbot_rsa_knowledge k
            WHERE json_query(k.metadata, '$.title') IN (
                SELECT DISTINCT title FROM top_chunks WHERE title IS NOT NULL
            )
        """

        sr = get_db_connection()
        results = await sr.execute_query_async(sql)

        if not results:
            logger.warning(f"No knowledge data found in DB for query: {query}")
            return "Hiện tại tôi chưa tìm thấy thông tin chính xác về nội dung này trong hệ thống kiến thức của Canifa. Bạn có thể liên hệ hotline 1800 6061 để được hỗ trợ trực tiếp."

        # 3. Phân nhóm nội dung theo title
        docs_by_title: dict[str, list[str]] = {}

        for res in results:
            try:
                meta_obj = (
                    json.loads(res.get("metadata", "{}"))
                    if isinstance(res.get("metadata"), str)
                    else res.get("metadata", {})
                )
                title = meta_obj.get("title") if isinstance(meta_obj, dict) else None
                content = res.get("content", "").strip()
                if title and content:
                    docs_by_title.setdefault(title, []).append(content)
            except Exception as e:
                logger.debug(f"Skip invalid row while grouping by title: {e}")

        if not docs_by_title:
            # Fallback: trả raw content nếu không parse được title
            knowledge_texts = [res.get("content", "") for res in results]
            return "\n\n---\n\n".join(knowledge_texts)

        # 4. Tổng hợp kết quả
        final_blocks = []
        for title, contents in docs_by_title.items():
            if contents:
                full_text = "\n\n".join(contents)
                final_blocks.append(f"=== TÀI LIỆU: {title} ===\n{full_text}")

        final_response = "\n\n".join(final_blocks)

        logger.info(f"Parent Document Retrieval complete. Loaded full context for {len(final_blocks)} topics.")
        return final_response

    except Exception as e:
        logger.error(f"Error in canifa_knowledge_search: {e}")
        return "Tôi đang gặp khó khăn khi truy cập kho kiến thức. Bạn muốn hỏi về sản phẩm gì khác không?"


canifa_knowledge_search.__doc__ = read_tool_prompt("brand_knowledge_tool") or canifa_knowledge_search.__doc__
