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:
    """
    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
        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. Tìm kiếm vector trên StarRocks
        sql = f"""
            WITH top_chunks AS (
                SELECT CAST(json_query(metadata, '$.title') AS VARCHAR) 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 CAST(json_query(k.metadata, '$.title') AS VARCHAR) 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:
            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."

        docs_by_title: dict[str, list[str]] = {}
        for res in results:
            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)

        final_blocks = []
        for title, contents in docs_by_title.items():
            full_text = "\n\n".join(contents)
            final_blocks.append(f"=== TÀI LIỆU: {title} ===\n{full_text}")

        return "\n\n".join(final_blocks)

    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."

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