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__)

# Constants
TOP_K_RESULTS = 10


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
        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. Query StarRocks lấy Top K kết quả phù hợp nhất
        sql = f"""
        SELECT
            content,
            metadata
        FROM shared_source.chatbot_rsa_knowledge
        ORDER BY approx_cosine_similarity(embedding, {v_str}) DESC
        LIMIT {TOP_K_RESULTS}
        """

        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. Tổng hợp kết quả
        knowledge_texts = [res.get("content", "") for res in results]
        final_response = "\n\n---\n\n".join(knowledge_texts)

        logger.info(f"✅ Found {len(results)} relevant knowledge chunks (top {TOP_K_RESULTS})")

        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__
