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 = 4


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ụ: tìm cửa hàng, hỏi chính sách, tra bảng size...). KHÔNG DÙNG ĐỂ TÌM SẢN PHẨM."
    )


@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 (Cửa hàng, chính sách, KHTT...).
    """
    logger.info(f"🔍 [Semantic Search] Brand Knowledge query: {query}")
    try:
        # 1. Tạo embedding cho câu hỏi (Mặc định 1536 chiều như bro yêu cầu)
        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__
