"""
CANIFA Data Retrieval Tool - Tối giản cho Agentic Workflow.
Hỗ trợ Hybrid Search: Semantic (Vector) + Metadata Filter.
"""

import asyncio
import json
import logging
import time
from decimal import Decimal

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

from agent.tools.product_search_helpers import build_starrocks_query
from common.embedding_service import create_embeddings_async
from common.starrocks_connection import get_db_connection

# from langsmith import traceable

logger = logging.getLogger(__name__)


class DecimalEncoder(json.JSONEncoder):
    """Xử lý kiểu Decimal từ Database khi convert sang JSON."""

    def default(self, obj):
        if isinstance(obj, Decimal):
            return float(obj)
        return super().default(obj)


class SearchItem(BaseModel):
    """
    Cấu trúc một mục tìm kiếm đơn lẻ trong Multi-Search.

    Lưu ý quan trọng về cách SINH QUERY:
    - Trường `query` KHÔNG phải câu hỏi thô của khách.
    - Phải là một đoạn text có cấu trúc giống hệt format trong cột `description_text_full` của DB,
      ví dụ (chỉ là 1 chuỗi duy nhất, nối các field bằng dấu chấm):

      product_name: Pack 3 đôi tất bé gái cổ thấp. master_color: Xanh da trời/ Blue.
      product_image_url: https://.... product_image_url_thumbnail: https://....
      product_web_url: https://.... description_text: ... material: ...
      material_group: Yarn - Sợi. gender_by_product: female. age_by_product: others.
      season: Year. style: Feminine. fitting: Slim. size_scale: 4/6.
      form_neckline: None. form_sleeve: None. product_line_vn: Tất.
      product_color_name: Blue Strip 449.

    - Khi khách chỉ nói “áo màu hồng”, hãy suy luận và sinh query dạng:

      product_name: Áo thun/áo sơ mi/áo ... màu hồng ... . master_color: Hồng/ Pink.
      product_image_url: None. product_image_url_thumbnail: None.
      product_web_url: None. description_text: ... (mô tả thêm nếu có).
      material: None. material_group: None. gender_by_product: ... (nếu đoán được).
      age_by_product: others. season: Year. style: ... (nếu đoán được).
      fitting: ... size_scale: None. form_neckline: None. form_sleeve: None.
      product_line_vn: Áo. product_color_name: Pink / Hồng (nếu hợp lý).

    - Nếu không suy luận được giá trị cho field nào thì để `None` hoặc bỏ trống phần text đó.
    """

    query: str = Field(
        ...,
        description=(
            "ĐOẠN TEXT CÓ CẤU TRÚC theo format của cột description_text_full trong DB, "
            "bao gồm các cặp key: product_name, master_color, product_image_url, "
            "product_image_url_thumbnail, product_web_url, description_text, material, "
            "material_group, gender_by_product, age_by_product, season, style, fitting, "
            "size_scale, form_neckline, form_sleeve, product_line_vn, product_color_name. "
            "Ví dụ: 'product_name: Pack 3 đôi tất bé gái cổ thấp. master_color: Xanh da trời/ Blue. "
            "product_image_url: https://.... product_web_url: https://.... description_text: ... "
            "material: None. material_group: Yarn - Sợi. gender_by_product: female. ...'"
        ),
    )
    magento_ref_code: str | None = Field(
        ..., description="Mã sản phẩm hoặc SKU (Ví dụ: 8TS24W001). CHỈ điền khi khách hỏi mã code cụ thể."
    )
    price_min: float | None = Field(..., description="Giá thấp nhất (VD: 100000)")
    price_max: float | None = Field(..., description="Giá cao nhất (VD: 500000)")
    action: str = Field(..., description="Hành động: 'search' (tìm kiếm) hoặc 'visual_search' (phân tích ảnh)")


class MultiSearchParams(BaseModel):
    """Tham số cho Parallel Multi-Search."""

    searches: list[SearchItem] = Field(..., description="Danh sách các truy vấn tìm kiếm chạy song song")



@tool(args_schema=MultiSearchParams)
# @traceable(run_type="tool", name="data_retrieval_tool")
async def data_retrieval_tool(searches: list[SearchItem]) -> str:
    """
    Siêu công cụ tìm kiếm sản phẩm CANIFA - Hỗ trợ Parallel Multi-Search (chạy song song nhiều truy vấn).

    Hướng dẫn dùng nhanh:
    - Trường 'query': mô tả chi tiết sản phẩm (tên, chất liệu, giới tính, màu sắc, phong cách, dịp sử dụng), không dùng câu hỏi thô.
    - Trường 'magento_ref_code': chỉ dùng khi khách hỏi mã sản phẩm/SKU cụ thể (vd: 8TS24W001).
    - Trường 'price_min' / 'price_max': dùng khi khách nói về khoảng giá (vd: dưới 500k, từ 200k đến 400k).
    """
    logger.info("data_retrieval_tool started, searches=%s", len(searches))
    try:
        # 0. Log input tổng quan (không log chi tiết dài)
        for idx, item in enumerate(searches):
            short_query = (item.query[:60] + "...") if item.query and len(item.query) > 60 else item.query
            logger.debug(
                "search[%s] query=%r, code=%r, price_min=%r, price_max=%r",
                idx,
                short_query,
                item.magento_ref_code,
                item.price_min,
                item.price_max,
            )


        queries_to_embed = [s.query for s in searches if s.query]
        all_vectors = []
        if queries_to_embed:
            logger.info("batch embedding %s queries", len(queries_to_embed))
            emb_batch_start = time.time()
            all_vectors = await create_embeddings_async(queries_to_embed)
            logger.info(
                "batch embedding done in %.2f ms",
                (time.time() - emb_batch_start) * 1000,
            )

        # 2. Get DB connection (singleton)
        db = get_db_connection()

        tasks = []
        vector_idx = 0
        for item in searches:
            current_vector = None
            if item.query:
                if vector_idx < len(all_vectors):
                    current_vector = all_vectors[vector_idx]
                    vector_idx += 1

            tasks.append(_execute_single_search(db, item, query_vector=current_vector))

        results = await asyncio.gather(*tasks)

        # 3. Tổng hợp kết quả
        combined_results = []
        for i, products in enumerate(results):
            combined_results.append(
                {
                    "search_index": i,
                    "search_criteria": searches[i].dict(exclude_none=True),
                    "count": len(products),
                    "products": products,
                }
            )

        logger.info("data_retrieval_tool finished, results=%s", len(combined_results))
        return json.dumps(
            {"status": "success", "results": combined_results},
            ensure_ascii=False,
            cls=DecimalEncoder,
        )

    except Exception as e:
        logger.exception("Error in Multi-Search data_retrieval_tool: %s", e)
        return json.dumps({"status": "error", "message": str(e)})


async def _execute_single_search(db, item: SearchItem, query_vector: list[float] | None = None) -> list[dict]:
    """Thực thi một search query đơn lẻ (Async)."""
    try:
        short_query = (item.query[:60] + "...") if item.query and len(item.query) > 60 else item.query
        logger.debug(
            "_execute_single_search started, query=%r, code=%r",
            short_query,
            item.magento_ref_code,
        )

        # Timer: build query (sử dụng vector đã có hoặc build mới)
        query_build_start = time.time()
        sql = await build_starrocks_query(item, query_vector=query_vector)
        query_build_time = (time.time() - query_build_start) * 1000  # Convert to ms
        logger.debug("SQL built, length=%s, build_time_ms=%.2f", len(sql), query_build_time)

        # Timer: execute DB query
        db_start = time.time()
        products = await db.execute_query_async(sql)
        db_time = (time.time() - db_start) * 1000  # Convert to ms
        logger.info(
            "_execute_single_search done, products=%s, build_ms=%.2f, db_ms=%.2f, total_ms=%.2f",
            len(products),
            query_build_time,
            db_time,
            query_build_time + db_time,
        )

        return _format_product_results(products)
    except Exception as e:
        logger.exception("Single search error for item %r: %s", item, e)
        return []


def _format_product_results(products: list[dict]) -> list[dict]:
    """Lọc và format kết quả trả về cho Agent."""
    max_items = 15
    formatted: list[dict] = []

    for p in products[:max_items]:
        formatted.append(
            {
                "internal_ref_code": p.get("internal_ref_code"),
                # Chuỗi text dài, đã bao gồm: product_name, master_color, image, web_url, material, style, ...
                "description_text": p.get("description_text_full"),
                "sale_price": p.get("sale_price"),
                "original_price": p.get("original_price"),
                "discount_amount": p.get("discount_amount"),
                "max_score": p.get("max_score"),
            }
        )

    return formatted
