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 - Chỉ cần lấy Top 2-3 chunk vì ta sẽ mở rộng lấy full dữ liệu của title đó
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
        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 chunks phù hợp nhất
        sql_search = f"""
                SELECT
                    content,
                    metadata
                FROM shared_source.chatbot_rsa_knowledge
                ORDER BY approx_cosine_similarity(embedding, {v_str}) DESC
                LIMIT {TOP_K_CHUNKS}
        """

        sr = get_db_connection()
        chunk_results = await sr.execute_query_async(sql_search)

        if not chunk_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. Trích xuất danh sách các 'title' duy nhất từ metadata của các chunk tìm được
        found_titles = set()
        for res in chunk_results:
            try:
                meta_str = res.get("metadata", "{}")
                meta_dict = json.loads(meta_str) if isinstance(meta_str, str) else meta_str
                title_value = meta_dict.get("title") if isinstance(meta_dict, dict) else None
                if title_value:
                    found_titles.add(title_value)
            except Exception as e:
                logger.error(f"Error parsing metadata: {e}")
                continue

        if not found_titles:
            # Fallback nếu không có title nào, trả về chunk như cũ
            knowledge_texts = [res.get("content", "") for res in chunk_results]
            return "\n\n---\n\n".join(knowledge_texts)

        logger.info(f"Semantic search found related titles: {list(found_titles)}")

        # 4. Gom TOÀN BỘ nội dung của những title đã tìm thấy
        # StarRocks JSON extract syntax: json_query(metadata, '$.title')
        safe_titles = [str(t).replace("'", "''") for t in found_titles]
        titles_sql_list = ", ".join([f"'{t}'" for t in safe_titles])

        sql_full_docs = f"""
            SELECT content, metadata
            FROM shared_source.chatbot_rsa_knowledge
            WHERE json_query(metadata, '$.title') IN ({titles_sql_list})
        """

        full_doc_results = await sr.execute_query_async(sql_full_docs)

        # 5. Phân nhóm nội dung theo title để context gọn gàng
        docs_by_title = {t: [] for t in found_titles}

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

        # 6. Tổng hợp kết quả
        final_blocks = []
        for title, contents in docs_by_title.items():
            if contents:
                # Ghép tất cả các chunk thuộc cùng một title
                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__
