"""
Lead Search Tool - LangChain @tool cho AI goi.
AI sinh keywords -> Tool search -> cascade -> Tra ket qua.

Architecture:
  CO DINH (luon giu): product_type + color + gender + price
  BIEN DOI (fallback):
    Tang 1: CO DINH + keywords LIKE description_text
    Tang 2: CO DINH + embedding cosine (vector column)
    Tang 3: CHI CO DINH (product_type + color + gender)
"""

import json
import logging
import time

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

from common.starrocks_connection import get_db_connection
from common.embedding_service import create_embedding_async
from common.pool_wrapper import get_pooled_connection_compat
from .product_mapping import SYNONYM_TO_DB, get_related_lines

try:
    from langfuse import get_client as get_langfuse
    _LANGFUSE_AVAILABLE = True
except ImportError:
    _LANGFUSE_AVAILABLE = False

def _lf_span(name: str, input_data: dict | None = None):
    """Create a Langfuse span if available, else return a no-op context manager."""
    if not _LANGFUSE_AVAILABLE:
        import contextlib
        return contextlib.nullcontext()
    try:
        lf = get_langfuse()
        if lf:
            return lf.start_as_current_observation(
                as_type="span",
                name=name,
                input=input_data or {},
            )
    except Exception:
        pass
    import contextlib
    return contextlib.nullcontext()


logger = logging.getLogger(__name__)

TABLE_NAME = "test_db.magento_product_dimension_with_text_embedding"
PG_TABLE   = "dashboard_canifa.ultra_descriptions"  # Postgres — can write/join

SELECT_COLUMNS = """
    internal_ref_code,
    magento_ref_code,
    product_color_code,
    product_name,
    master_color,
    product_color_name,
    product_image_url_thumbnail,
    product_web_url,
    sale_price,
    original_price,
    COALESCE(discount_amount, 0) AS discount_amount,
    ROUND(CASE WHEN original_price > 0
        THEN ((original_price - sale_price) / original_price * 100)
        ELSE 0 END, 0) AS discount_percent,
    age_by_product,
    gender_by_product,
    product_line_vn,
    COALESCE(quantity_sold, 0) AS quantity_sold,
    COALESCE(is_new_product, 0) AS is_new_product,
    size_scale,
    description_text,
    suggest_items,
    similar_items
"""


# ======================================================
# Pydantic Schema - AI se fill vao day
# ======================================================
class LeadSearchInput(BaseModel):
    """Schema cho hybrid search: keywords (LIKE) + embedding (cosine) fallback."""

    model_config = {"extra": "forbid"}

    # KEYWORDS - Rut gon intent tu cau hoi user
    keywords: list[str] = Field(
        default=[],
        description=(
            "Rut gon y dinh khach thanh cac CUM TU NGAN de search MO TA SAN PHAM. "
            "TUYET DOI CAM dua MAU SAC vao keywords — mau da co truong master_color rieng. "
            "Chi dua: chat lieu (cotton, interlock, denim), phong cach (oversize, ong rong), "
            "dip (dam cuoi, di bien), ky nang che dau nhuoc diem (che bung, che dui)... "
            "VD: khach noi 'ao do chat lieu interlock' -> keywords=['interlock'], master_color='do'. "
            "VD: khach noi 'di dam cuoi troi 40 do beo bung' -> keywords=['dam cuoi', 'thoang mat', 'che bung']. "
            "VD: khach noi 'cotton ong rong' -> keywords=['cotton', 'ong rong']. "
            "VD: khach noi 'goi y do dep' -> keywords=[]. "
            "Toi da 4. De [] neu khach hoi chung chung."
        ),
    )

    # CO DINH - Luon giu trong moi tang
    product_line_vn: list[str] = Field(
        default=[],
        description=(
            "Danh muc san pham chuẩn (VD: 'Áo phông', 'Áo Polo', 'Váy liền'...). "
            "LUON UU TIEN dien neu xac dinh duoc loai san pham khach muon. "
            "Dung de filter chinh xac theo cot product_line_vn trong DB."
        ),
    )
    gender_by_product: str | None = Field(
        default=None,
        description="Gioi tinh. GIA TRI HOP LE: women, men, boy, girl, unisex, newborn",
    )
    age_by_product: str | None = Field(
        default=None,
        description="Do tuoi. GIA TRI HOP LE: adult, kid, others",
    )
    master_color: str | None = Field(
        default=None,
        description=(
            "Mau sac. Ghi CHINH XAC tu khach noi: 'trang', 'den', 'do'... "
            "VD: 'ao co do sao vang' -> color='do'."
        ),
    )
    price_min: int | None = Field(default=None, description="Gia thap nhat (VND)")
    price_max: int | None = Field(default=None, description="Gia cao nhat (VND)")
    discount_min: int | None = Field(
        default=None,
        description="% giam gia TOI THIEU. VD: 'giam 50%' -> 50. 'hang sale' -> 1.",
    )
    discount_max: int | None = Field(
        default=None,
        description="% giam gia TOI DA. VD: 'giam duoi 30%' -> 30.",
    )
    discovery_mode: str | None = Field(
        default=None,
        description="'new' = hang moi, 'best_seller' = ban chay. Chi dung khi khach noi ro.",
    )
    size: str | None = Field(
        default=None,
        description=(
            "Size san pham khach yeu cau. "
            "Nguoi lon: XS, S, M, L, XL, XXL, 3XL, 4XL. "
            "Tre em (chieu cao cm): 80, 86, 92, 98, 104, 110, 116, 122, 128, 134, 140, 152, 164. "
            "Tre so sinh (thang): 1/3, 4/6, 7/9, 10/12, 13/14. "
            "VD: 'ao size L' -> size='L'. 'quan be 110' -> size='110'. "
            "Chi dien khi khach NOI RO size. Neu khong nhac -> None."
        ),
    )
    magento_ref_code: str | None = Field(
        default=None,
        description="Ma SKU cu the. Chi dien khi khach cung cap.",
    )
    reasoning: str | None = Field(
        default=None,
        description=(
            "SUY LUAN: Giai thich ngan gon TAI SAO ban chon params nay. "
            "VD: '2/9 = Quoc Khanh -> ao do di choi', 'khach muon di dam cuoi troi nong -> thoang mat, lich su'. "
            "Viet 1-2 cau bang tieng Viet."
        ),
    )


# ======================================================
# SQL Builder - Hard Filters
# ======================================================
def _build_fixed_clauses(req: LeadSearchInput, params: list) -> list[str]:
    """Build các filter CỐ ĐỊNH (luôn giữ trong mọi tầng)."""
    clauses = []

    # Product Line VN
    if req.product_line_vn:
        lines = []
        for line in req.product_line_vn:
            if not line: continue
            # Resolve synonyms (e.g. "áo polo" -> "Áo Polo")
            from .product_mapping import resolve_product_line
            resolved = resolve_product_line(line)
            for r in resolved:
                # Expand related lines (e.g. "Áo lót" -> ["Áo lót", "Áo bra active"])
                expanded = get_related_lines(r)
                lines.extend(expanded)
        
        if lines:
            placeholders = ", ".join(["%s"] * len(lines))
            params.extend(lines)
            clauses.append(f"product_line_vn IN ({placeholders})")

    # Gender
    if req.gender_by_product:
        gender_lower = req.gender_by_product.lower().strip()
        genders_to_search = []
        
        if gender_lower in ("women", "nu", "female", "nữ"):
            genders_to_search = ["female", "women", "nu", "nữ", "unisex"]
        elif gender_lower in ("men", "nam", "male"):
            genders_to_search = ["male", "men", "nam", "unisex"]
        elif gender_lower in ("boy", "bé trai", "be trai"):
            genders_to_search = ["boy", "bé trai", "be trai", "unisex"]
        elif gender_lower in ("girl", "bé gái", "be gai"):
            genders_to_search = ["girl", "bé gái", "be gai", "unisex"]
        else:
            genders_to_search = [gender_lower, "unisex"]

        placeholders = ", ".join(["%s"] * len(genders_to_search))
        params.extend(genders_to_search)
        clauses.append(f"gender_by_product IN ({placeholders})")

    # Age
    if req.age_by_product:
        age_lower = req.age_by_product.lower().strip()
        if age_lower in ("baby", "newborn"):
            params.append("kid")
        else:
            params.append(age_lower)
        clauses.append("age_by_product = %s")

    # Color — StarRocks LOWER() khong handle Vietnamese diacritics
    # → match nhieu variant: tieng Viet (goc + khong dau) + tieng Anh
    if req.master_color:
        c_raw = req.master_color.strip()

        # Map mau Viet -> EN. Bao gom ca dang co dau va khong dau
        # Vi "d-stroke" (U+0111) khong decompose qua NFD nen phai map thu cong
        COLOR_EN_MAP = {
            # Khong dau
            "do": "red", "trang": "white", "den": "black", "xanh": "blue",
            "hong": "pink", "tim": "purple", "cam": "orange", "vang": "yellow",
            "nau": "brown", "xam": "gray", "be": "beige", "kem": "cream",
            "xanh la": "green", "bac": "silver", "vang dong": "gold",
            # Co dau (d-stroke variants)
            "đo": "red",     # đo
            "đỏ": "red", # đỏ (capital Đ lower = đ, o + hook above = ỏ)
            "trắng": "white", # trắng
            "đen": "black",  # đen
            "xanh lá": "green",  # xanh lá
            "hồng": "pink",  # hồng
            "tím": "purple", # tím
            "nâu": "brown",  # nâu
            "xám": "gray",   # xám
            "vàng": "yellow", # vàng
            "bạc": "silver", # bạc
        }
        c_key = c_raw.lower().strip()
        en_color = COLOR_EN_MAP.get(c_key)
        if not en_color:
            # fallback: strip combining marks only (works for o/a/e etc, not d-stroke)
            import unicodedata
            c_no_acc = "".join(
                ch for ch in unicodedata.normalize("NFD", c_key)
                if unicodedata.category(ch) != "Mn"
            ).strip()
            en_color = COLOR_EN_MAP.get(c_no_acc)
        c_no_accent = c_key  # used below for non-LOWER match

        color_parts = []
        # 1. Match chinh xac khong qua LOWER (tranh loi diacritics)
        color_parts.append("(master_color LIKE %s OR product_color_name LIKE %s)")
        params.append(f"%{c_raw}%")
        params.append(f"%{c_raw}%")
        # 2. Match khong dau (VD: "do" -> match "Do", "do", "DO")
        if c_no_accent != c_raw.lower():
            color_parts.append("(LOWER(master_color) LIKE %s OR LOWER(product_color_name) LIKE %s)")
            params.append(f"%{c_no_accent}%")
            params.append(f"%{c_no_accent}%")
        # 3. Match tieng Anh (VD: "do" -> "red")
        if en_color:
            color_parts.append("(LOWER(master_color) LIKE %s OR LOWER(product_color_name) LIKE %s)")
            params.append(f"%{en_color}%")
            params.append(f"%{en_color}%")

        clauses.append("(" + " OR ".join(color_parts) + ")")

    # Price
    if req.price_min is not None:
        params.append(req.price_min)
        clauses.append("sale_price >= %s")
    if req.price_max is not None:
        params.append(req.price_max)
        clauses.append("sale_price <= %s")

    # Discount
    if req.discount_min is not None:
        params.append(req.discount_min)
        clauses.append("(original_price > 0 AND ((original_price - sale_price) / original_price * 100) >= %s)")
    if req.discount_max is not None:
        params.append(req.discount_max)
        clauses.append("(original_price > 0 AND ((original_price - sale_price) / original_price * 100) <= %s)")

    # Size
    if req.size:
        s = req.size.strip().upper()
        params.append(s)
        clauses.append("FIND_IN_SET(%s, REPLACE(size_scale, '|', ',')) > 0")

    # Discovery mode
    if req.discovery_mode:
        mode = req.discovery_mode.lower().strip()
        if mode == "new":
            clauses.append("is_new_product = 1")
        elif mode == "best_seller":
            clauses.append("quantity_sold > 0")

    return clauses


# ======================================================
# Tier 1: Keywords LIKE search on description_text
# ======================================================
def _build_keyword_clause(keywords: list[str], params: list) -> str | None:
    """Build menh de LIKE cho tung keyword, AND voi nhau de tang do chinh xac."""
    if not keywords:
        return None
    parts = []
    for kw in keywords[:4]:
        kw = kw.strip()
        if not kw:
            continue
        params.append(f"%{kw}%")
        params.append(f"%{kw}%")
        # Search ca description_text (mien ngan) va description_text_full (day du)
        parts.append("(LOWER(description_text) LIKE LOWER(%s) OR LOWER(description_text_full) LIKE LOWER(%s))")
    if not parts:
        return None
    # Dùng OR thay vì AND để tăng độ phủ (Recall) cho tìm kiếm từ khóa
    return "(" + " OR ".join(parts) + ")"


def _build_exclusion_clauses(keywords: list[str], params: list) -> list[str]:
    """Build các filter LOẠI TRỪ (Negative filters) dựa trên từ khóa."""
    clauses = []
    kws_str = " ".join(keywords).lower()

    # 1. Mùa đông/Lạnh -> Loại bỏ đồ hè hở hang
    if any(k in kws_str for k in ["đông", "lạnh", "winter", "tuyết", "giữ ấm"]):
        # Loại trừ Quần soóc, Áo ba lỗ
        clauses.append("product_line_vn NOT IN (%s, %s)")
        params.extend(["Quần soóc", "Áo ba lỗ"])

    # 2. Đi làm/Công sở -> Loại bỏ đồ hoạt hình/manga nếu có thể
    if any(k in kws_str for k in ["đi làm", "công sở", "văn phòng", "office"]):
        # Loại trừ các mô tả chứa cartoon/manga/anime
        # Dùng LIKE NOT để lọc bớt
        forbidden = ["cartoon", "hoạt hình", "manga", "anime", "demon slayer", "naruto", "disney", "marvel"]
        for f in forbidden:
            clauses.append("LOWER(description_text) NOT LIKE %s")
            params.append(f"%{f}%")

    return clauses


def _build_full_query(fixed_clauses: list[str], keyword_clause: str | None, exclusion_clauses: list[str] = None) -> str:
    """Lắp ráp câu SELECT hoàn chỉnh từ fixed + keyword + exclusion clauses."""
    all_clauses = fixed_clauses[:]
    if keyword_clause:
        all_clauses.append(keyword_clause)
    if exclusion_clauses:
        all_clauses.extend(exclusion_clauses)

    where_sql = ""
    if all_clauses:
        where_sql = "WHERE " + " AND ".join(all_clauses)

    return f"""
    SELECT
        internal_ref_code,
        magento_ref_code,
        product_color_code,
        product_name,
        master_color,
        product_color_name,
        product_image_url_thumbnail,
        product_web_url,
        sale_price,
        original_price,
        COALESCE(discount_amount, 0) AS discount_amount,
        ROUND(CASE WHEN original_price > 0
            THEN ((original_price - sale_price) / original_price * 100)
            ELSE 0 END, 0) AS discount_percent,
        age_by_product,
        gender_by_product,
        product_line_vn,
        COALESCE(quantity_sold, 0) AS quantity_sold,
        COALESCE(is_new_product, 0) AS is_new_product,
        size_scale,
        description_text,
        suggest_items,
        similar_items
    FROM {TABLE_NAME}
    {where_sql}
    ORDER BY quantity_sold DESC NULLS LAST
    LIMIT 20
    """


# ======================================================
# Tier 2: Embedding semantic search (HNSW ANN)
# ======================================================
async def _embedding_search_with_vec(
    query_vec: list,
    query_str: str,
    fixed_clauses: list[str],
    fixed_params: list,
    db,
    exclusion_clauses: list[str] = None,
) -> list:
    """
    Nhận vector đã compute -> approx_cosine_similarity trên StarRocks.
    CTE: vector_matches(top200) -> filtered(hard) -> GROUP BY dedup -> top 20.
    """
    t0 = time.time()
    vec_literal = "[" + ",".join(str(v) for v in query_vec) + "]"

    post_filter_clauses = fixed_clauses[:]
    if exclusion_clauses:
        post_filter_clauses.extend(exclusion_clauses)

    post_filter_where = ""
    if post_filter_clauses:
        post_filter_where = " WHERE " + " AND ".join(post_filter_clauses)

    sql = f"""
    WITH vector_matches AS (
        SELECT /*+ SET_VAR(ann_params='{{"ef_search":256}}') */
            internal_ref_code,
            magento_ref_code,
            product_color_code,
            product_name,
            master_color,
            product_color_name,
            product_image_url_thumbnail,
            product_web_url,
            sale_price,
            original_price,
            COALESCE(discount_amount, 0) AS discount_amount,
            ROUND(CASE WHEN original_price > 0
                THEN ((original_price - sale_price) / original_price * 100)
                ELSE 0 END, 0) AS discount_percent,
            age_by_product,
            gender_by_product,
            product_line_vn,
            COALESCE(quantity_sold, 0) AS quantity_sold,
            COALESCE(is_new_product, 0) AS is_new_product,
            size_scale,
            description_text,
            suggest_items,
            similar_items,
            approx_cosine_similarity(vector, {vec_literal}) AS similarity_score
        FROM {TABLE_NAME}
        ORDER BY similarity_score DESC
        LIMIT 200
    ),
    filtered_matches AS (
        SELECT * FROM vector_matches
        {post_filter_where}
        ORDER BY similarity_score DESC
        LIMIT 100
    )
    SELECT
        internal_ref_code,
        MAX_BY(magento_ref_code, similarity_score) AS magento_ref_code,
        product_color_code,
        MAX_BY(product_name, similarity_score) AS product_name,
        MAX_BY(master_color, similarity_score) AS master_color,
        MAX_BY(product_color_name, similarity_score) AS product_color_name,
        MAX_BY(product_image_url_thumbnail, similarity_score) AS product_image_url_thumbnail,
        MAX_BY(product_web_url, similarity_score) AS product_web_url,
        MAX_BY(sale_price, similarity_score) AS sale_price,
        MAX_BY(original_price, similarity_score) AS original_price,
        MAX_BY(discount_amount, similarity_score) AS discount_amount,
        MAX_BY(ROUND(CASE WHEN original_price > 0
            THEN ((original_price - sale_price) / original_price * 100)
            ELSE 0 END, 0), similarity_score) AS discount_percent,
        MAX_BY(age_by_product, similarity_score) AS age_by_product,
        MAX_BY(gender_by_product, similarity_score) AS gender_by_product,
        MAX_BY(product_line_vn, similarity_score) AS product_line_vn,
        MAX_BY(quantity_sold, similarity_score) AS quantity_sold,
        MAX_BY(is_new_product, similarity_score) AS is_new_product,
        MAX_BY(size_scale, similarity_score) AS size_scale,
        MAX_BY(description_text, similarity_score) AS description_text,
        ANY_VALUE(suggest_items) AS suggest_items,
        ANY_VALUE(similar_items) AS similar_items,
        MAX(similarity_score) AS _score
    FROM filtered_matches
    GROUP BY product_color_code, internal_ref_code
    ORDER BY _score DESC
    LIMIT 20
    """

    products = await db.execute_query_async(sql, params=tuple(fixed_params) if fixed_params else None)
    search_ms = round((time.time() - t0) * 1000, 1)
    logger.info("[EMBEDDING] query='%s' -> %d results | search=%.0fms", query_str[:60], len(products), search_ms)
    return products


# ======================================================
# Cascading Search
# ======================================================
async def _price_relaxed_search(req: LeadSearchInput, db, multiplier: float) -> list:
    """Tim voi gia noi (nhan multiplier), giu nguyen product_line + gender."""
    if req.price_max is None:
        return []
    saved_max = req.price_max
    saved_min = req.price_min
    req.price_max = int(req.price_max * multiplier)
    req.price_min = None
    params: list = []
    fixed = _build_fixed_clauses(req, params)
    sql = _build_full_query(fixed, None)
    products = await db.execute_query_async(sql, params=tuple(params))
    req.price_max = saved_max
    req.price_min = saved_min
    return products


async def _cascading_search(
    req: LeadSearchInput,
    db,
    raw_message: str = "",
) -> tuple[list, int, str | None, dict]:
    """
    Cascading search:
      Tang 1: CO DINH + keywords LIKE description_text
      Tang 2: CO DINH + embedding cosine (keywords joined -> embed)
      Tang 3: CHI CO DINH (product_line + gender + price)
      Tang 4: Drop gender
      Tang 5-6: Price relaxation (1.5x, 2x)
      Tang 7: Bo price hoan toan

    Returns: (products, tier_matched, fallback_message | None, search_details)
    """
    search_details: dict = {
        "keyword_query": None,
        "keyword_time_ms": None,
        "keyword_results": None,
        "embed_query": None,
        "openai_embed_ms": None,
        "starrocks_search_ms": None,
        "embed_total_ms": None,
        "embed_results": None,
    }
    _cascade_meta = {
        "product_line": req.product_line_vn,
        "gender": req.gender_by_product,
        "age": req.age_by_product,
        "color": req.master_color,
        "price_min": req.price_min,
        "price_max": req.price_max,
        "keywords": req.keywords,
        "raw_message": raw_message,
    }
    try:
        _lf = get_langfuse() if _LANGFUSE_AVAILABLE else None
    except Exception:
        _lf = None

    # ---- Prep Negative Filters (Exclusions) ----
    ex_params = []
    exclusions = _build_exclusion_clauses(req.keywords, ex_params)

    # ---- Tang 1: CO DINH + KEYWORDS LIKE ----
    if req.keywords:
        t1 = time.time()
        params: list = []
        fixed = _build_fixed_clauses(req, params)
        search = _build_keyword_clause(req.keywords, params)
        if search:
            # Combine params: fixed + search + exclusions
            all_params = params + ex_params
            sql = _build_full_query(fixed, search, exclusions)
            products = await db.execute_query_async(sql, params=tuple(all_params))
            kw_ms = round((time.time() - t1) * 1000, 1)
            search_details["keyword_query"] = req.keywords
            search_details["keyword_time_ms"] = kw_ms
            search_details["keyword_results"] = len(products)
            logger.info(
                "[SEARCH][TIER-1 KEYWORD] query=%s | filters={line=%s gender=%s color=%s price=%s~%s} -> %d results | %.0fms",
                req.keywords, req.product_line_vn, req.gender_by_product,
                req.master_color, req.price_min, req.price_max,
                len(products), kw_ms,
            )
            if _lf:
                try:
                    with _lf.start_as_current_observation(
                        as_type="span", name="search:tier1_keyword",
                        input={"keywords": req.keywords, "filters": _cascade_meta},
                        metadata={"sql_clauses_count": len(fixed)},
                    ) as _sp1:
                        _sp1.update(output={"results": len(products), "latency_ms": kw_ms, "hit": bool(products)})
                except Exception:
                    pass
            if products:
                logger.info("[SEARCH] => HIT Tier 1 (keyword LIKE) | %d results", len(products))
                return products, 1, None, search_details
            logger.info("[SEARCH] => MISS Tier 1 -> fallback embedding")

    # ---- Tang 2: CO DINH + EMBEDDING cosine ----
    embed_query = " ".join(req.keywords) if req.keywords else raw_message
    if embed_query.strip():
        t2 = time.time()

        # Step 2a: OpenAI Embedding
        query_vec = await create_embedding_async(embed_query.strip())
        openai_ms = round((time.time() - t2) * 1000, 1)

        if query_vec and len(query_vec) == 1536:
            # Step 2b: StarRocks cosine search
            t2b = time.time()
            params = []
            fixed = _build_fixed_clauses(req, params)
            # Combine params for embedding search
            all_params = params + ex_params
            products = await _embedding_search_with_vec(query_vec, embed_query, fixed, all_params, db, exclusions)
            sr_ms = round((time.time() - t2b) * 1000, 1)
            embed_total = round((time.time() - t2) * 1000, 1)

            search_details["embed_query"] = embed_query.strip()
            search_details["openai_embed_ms"] = openai_ms
            search_details["starrocks_search_ms"] = sr_ms
            search_details["embed_total_ms"] = embed_total
            search_details["embed_results"] = len(products)

            logger.info(
                "[SEARCH][TIER-2 SEMANTIC] query='%s' | filters={line=%s gender=%s} -> %d results | openai=%.0fms sr=%.0fms total=%.0fms",
                embed_query[:60], req.product_line_vn, req.gender_by_product,
                len(products), openai_ms, sr_ms, embed_total,
            )
            if _lf:
                try:
                    with _lf.start_as_current_observation(
                        as_type="span", name="search:tier2_semantic",
                        input={"embed_query": embed_query, "filters": _cascade_meta},
                        metadata={"openai_embed_ms": openai_ms, "starrocks_search_ms": sr_ms},
                    ) as _sp2:
                        _sp2.update(output={"results": len(products), "total_ms": embed_total, "hit": bool(products)})
                except Exception:
                    pass
            if products:
                logger.info("[SEARCH] => HIT Tier 2 (semantic embedding) | %d results", len(products))
                return products, 2, None, search_details
        else:
            search_details["openai_embed_ms"] = openai_ms
            logger.warning("[SEARCH] => Embedding FAILED for: %s", embed_query[:60])

    # ---- Tang 3: CHI CO DINH ----
    params = []
    fixed = _build_fixed_clauses(req, params)
    
    # Neu khong co filter cung (line, gender, color, price) thi Tier 3 KHONG tra gi ca 
    # de tranh tra ve best-seller toan cong ty khong lien quan.
    has_hard_filter = bool(req.product_line_vn or req.gender_by_product or req.master_color or req.price_max)
    
    if not has_hard_filter:
        logger.info("[SEARCH][TIER-3] SKIPPED - no hard filters")
        products = []
    else:
        all_params = params + ex_params
        sql = _build_full_query(fixed, None, exclusions)
        products = await db.execute_query_async(sql, params=tuple(all_params))
        logger.info(
            "[SEARCH][TIER-3 HARD-FILTERS] filters={line=%s gender=%s age=%s color=%s price=%s~%s disc=%s~%s size=%s mode=%s} -> %d results",
            req.product_line_vn, req.gender_by_product, req.age_by_product,
            req.master_color, req.price_min, req.price_max,
            req.discount_min, req.discount_max, req.size, req.discovery_mode,
            len(products),
        )
    if products:
        logger.info("[SEARCH] => HIT Tier 3 (hard filters only) | %d results", len(products))
        return products, 3, None, search_details
    logger.info("[SEARCH] => MISS Tier 3 -> try drop gender")

    # ---- Tang 4a: DROP GENDER ----
    if req.product_line_vn and req.gender_by_product:
        logger.info("[SEARCH][TIER-4a DROP-GENDER] dropping gender='%s' ...", req.gender_by_product)
        saved_gender = req.gender_by_product
        req.gender_by_product = None
        params = []
        fixed = _build_fixed_clauses(req, params)
        sql = _build_full_query(fixed, None)
        products = await db.execute_query_async(sql, params=tuple(params))
        req.gender_by_product = saved_gender
        if products:
            logger.info("[SEARCH] => HIT Tier 4a (drop gender) | %d results", len(products))
            return products, 4, None, search_details
        logger.info("[SEARCH] => MISS Tier 4a -> try drop color")

    # ---- Tang 4b: DROP COLOR ----
    if req.product_line_vn and req.master_color:
        logger.info("[SEARCH][TIER-4b DROP-COLOR] dropping color='%s' ...", req.master_color)
        saved_color = req.master_color
        req.master_color = None
        params = []
        fixed = _build_fixed_clauses(req, params)

        # Khi drop color: loai tat ca keywords lien quan mau sac ra khoi embed query
        # chi giu lai keywords ve chuc nang / chat lieu / phong cach
        _color_terms = {saved_color.lower()} if saved_color else set()
        non_color_kws = [
            kw for kw in (req.keywords or [])
            if saved_color and saved_color.lower() not in kw.lower()
            and not any(c in kw.lower() for c in ("mau", "màu", "color"))
        ]
        embed_query_color = " ".join(non_color_kws) if non_color_kws else ""
        logger.info(
            "[SEARCH][TIER-4b] embed_query (no color): '%s' (original kws=%s)",
            embed_query_color, req.keywords,
        )

        if embed_query_color.strip():
            query_vec_color = await create_embedding_async(embed_query_color.strip())
            if query_vec_color and len(query_vec_color) == 1536:
                products = await _embedding_search_with_vec(query_vec_color, embed_query_color, fixed, params, db)
                logger.info(
                    "[SEARCH][TIER-4b DROP-COLOR EMBED] query='%s' -> %d results",
                    embed_query_color[:50], len(products),
                )
                if products:
                    req.master_color = saved_color
                    logger.info("[SEARCH] => HIT Tier 4b (drop color + semantic) | %d results", len(products))
                    return products, 4, None, search_details

        # fallback: chi hard filter, khong co color, khong co keywords
        sql = _build_full_query(fixed, None)
        products = await db.execute_query_async(sql, params=tuple(params))
        req.master_color = saved_color
        if products:
            logger.info("[SEARCH] => HIT Tier 4b (drop color + hard filter) | %d results", len(products))
            return products, 4, None, search_details
        logger.info("[SEARCH] => MISS Tier 4b -> price relaxation")

    # ---- Tang 5-7: PRICE RELAXATION ----
    if req.price_max is not None and req.product_line_vn:
        original_max = req.price_max
        product_label = "/".join(req.product_line_vn)

        for multiplier, tier, label in [(1.5, 5, "1.5x"), (2.0, 6, "2x")]:
            new_max = int(original_max * multiplier)
            logger.info("[SEARCH][TIER-%d PRICE-RELAX %s] original_max=%d -> new_max=%d", tier, label, original_max, new_max)
            products = await _price_relaxed_search(req, db, multiplier)
            if products:
                cheapest_p = min(products, key=lambda p: float(p.get("sale_price") or 999_999_999))
                cheapest = float(cheapest_p.get("sale_price") or 0)
                cheapest_name = cheapest_p.get("product_name", "")
                diff = int(cheapest - original_max)
                msg = (
                    f"FALLBACK_PRICE: Khong co {product_label} duoi {original_max:,.0f}d. "
                    f"Mau re nhat: \"{cheapest_name}\" gia {cheapest:,.0f}d "
                    f"(chi them {diff:,}d so voi budget). "
                    f"Tong co {len(products)} mau trong tam {new_max:,.0f}d. "
                    f"-> HAY goi y SP nay cho khach va noi kheo: them chut xiu la co mau rat dep!"
                )
                return products, tier, msg, search_details

        # Tang 7: bo price hoan toan
        logger.info("[SEARCH][TIER-7 PRICE-UNLIMITED] dropping price entirely (original_max=%d)", original_max)
        saved_max = req.price_max
        saved_min = req.price_min
        req.price_max = None
        req.price_min = None
        params = []
        fixed = _build_fixed_clauses(req, params)
        sql = _build_full_query(fixed, None)
        products = await db.execute_query_async(sql, params=tuple(params))
        req.price_max = saved_max
        req.price_min = saved_min
        if products:
            cheapest_p = min(products, key=lambda p: float(p.get("sale_price") or 999_999_999))
            cheapest = float(cheapest_p.get("sale_price") or 0)
            cheapest_name = cheapest_p.get("product_name", "")
            msg = (
                f"FALLBACK_PRICE: Khong co {product_label} duoi {original_max:,.0f}d. "
                f"Mau re nhat hien co: \"{cheapest_name}\" gia {cheapest:,.0f}d. "
                f"Tong co {len(products)} mau. "
                f"-> HAY goi y SP nay va noi: budget them chut la co mau rat dang!"
            )
            return products, 7, msg, search_details

    return [], 3, None, search_details


# ======================================================
# Postgres Enrichment — ai_description + ai_matches
# ======================================================
def _enrich_from_postgres(products: list[dict]) -> dict[str, dict]:
    """
    Batch-fetch extra columns from Postgres ultra_descriptions
    for a list of formatted products.
    Returns: {magento_ref_code: {ai_description, ai_matches, ...}}
    """
    codes = [p.get("sku") for p in products if p.get("sku")]
    if not codes:
        return {}

    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur  = conn.cursor()
        ph   = ", ".join(["%s"] * len(codes))
        cur.execute(
            f"""
            SELECT magento_ref_code,
                   description_data,
                   ai_matches
            FROM {PG_TABLE}
            WHERE magento_ref_code IN ({ph})
            """,
            tuple(codes),
        )
        rows = cur.fetchall()
        cur.close()
    except Exception as e:
        logger.warning("[LEAD] Postgres enrich error: %s", e)
        return {}
    finally:
        if conn:
            conn.close()

    enrich_map: dict[str, dict] = {}
    for code, dd_raw, am_raw in rows:
        # Parse description_data
        if isinstance(dd_raw, str):
            try:
                dd_raw = json.loads(dd_raw)
            except Exception:
                dd_raw = {}
        dd = dd_raw if isinstance(dd_raw, dict) else {}

        # Parse ai_matches
        if isinstance(am_raw, str):
            try:
                am_raw = json.loads(am_raw)
            except Exception:
                am_raw = {}
        am = am_raw if isinstance(am_raw, dict) else {}

        # Build a compact ai_description snippet for the LLM
        parts = []
        if dd.get("tagline"):      parts.append(dd["tagline"])
        if dd.get("phong_cach"):   parts.append(f"Phong cách: {dd['phong_cach']}")
        if dd.get("dip_mac"):      parts.append(f"Dịp: {dd['dip_mac']}")
        if dd.get("ly_do_mua"):    parts.append(f"Điểm nổi bật: {dd['ly_do_mua']}")
        if dd.get("chat_lieu"):    parts.append(f"Chất liệu: {dd['chat_lieu']}")

        enrich_map[code] = {
            "ai_description": " | ".join(parts) if parts else None,
            "ai_matches":     am or None,
        }
    return enrich_map


async def _format_products(products: list, db) -> list[dict]:
    """Format output cho AI doc, dong thoi TU DONG TRA CUU thong tin suggest/similar items."""
    formatted = []

    all_related_skus: set = set()
    raw_parsed_map: dict = {}

    def parse_item_list(raw_val) -> list[str]:
        if not raw_val:
            return []
        if isinstance(raw_val, list):
            return [str(x) for x in raw_val[:3]]
        if isinstance(raw_val, str):
            try:
                parsed = json.loads(raw_val)
                if isinstance(parsed, list):
                    return [str(x) for x in parsed[:3]]
            except Exception:
                pass
        return []

    for p in products[:20]:
        sku = p.get("magento_ref_code", "")
        sug_skus = parse_item_list(p.get("suggest_items"))[:2]
        sim_skus = parse_item_list(p.get("similar_items"))[:2]
        raw_parsed_map[sku] = {"suggest": sug_skus, "similar": sim_skus}
        all_related_skus.update(sug_skus)
        all_related_skus.update(sim_skus)

    sku_lookup: dict = {}
    if all_related_skus:
        phs = ", ".join(["%s"] * len(all_related_skus))
        q = f"SELECT magento_ref_code, product_name, sale_price FROM {TABLE_NAME} WHERE magento_ref_code IN ({phs})"
        related_products = await db.execute_query_async(q, params=tuple(all_related_skus))
        for rp in related_products:
            rpsku = rp.get("magento_ref_code")
            if rpsku:
                sku_lookup[rpsku] = {
                    "sku": rpsku,
                    "name": rp.get("product_name"),
                    "price": int(rp.get("sale_price") or 0),
                }

    for p in products[:20]:
        sale = float(p.get("sale_price") or 0)
        orig = float(p.get("original_price") or 0)
        has_discount = sale < orig and orig > 0
        disc_pct = int(p.get("discount_percent") or 0)
        desc = (p.get("description_text") or "").strip()
        desc_short = desc[:300] + "..." if len(desc) > 300 else desc

        main_sku = p.get("magento_ref_code", "")
        parsed_rel = raw_parsed_map.get(main_sku, {"suggest": [], "similar": []})
        rich_suggest = [sku_lookup[s] for s in parsed_rel["suggest"] if s in sku_lookup]
        rich_similar = [sku_lookup[s] for s in parsed_rel["similar"] if s in sku_lookup]

        item = {
            "sku": main_sku,
            "name": p.get("product_name", ""),
            "price": int(sale),
            "original_price": int(orig),
            "discount": f"-{disc_pct}%" if has_discount else None,
            "color": p.get("master_color", ""),
            "gender": p.get("gender_by_product", ""),
            "product_line": p.get("product_line_vn", ""),
            "image": p.get("product_image_url_thumbnail", ""),
            "url": p.get("product_web_url", ""),
            "sizes": p.get("size_scale", ""),
            "description": desc_short,
            "suggest_items": rich_suggest,
            "similar_items": rich_similar,
        }
        # Embedding score neu co
        if "_score" in p:
            item["relevance_score"] = round(float(p["_score"]), 4)
        formatted.append(item)

    # ── Enrich from Postgres (ai_description + ai_matches) ──
    try:
        pg_data = _enrich_from_postgres(formatted)
        for item in formatted:
            extra = pg_data.get(item.get("sku"), {})
            if extra.get("ai_description"):
                item["ai_description"] = extra["ai_description"]
            if extra.get("ai_matches"):
                item["ai_matches"] = extra["ai_matches"]
    except Exception as e:
        logger.warning("[LEAD] Postgres enrich merge error: %s", e)

    return formatted


# ======================================================
# SKU Lookup (dac biet, bypass cascade)
# ======================================================
def _build_sku_query(code: str) -> tuple[str, list]:
    """Build query tim theo ma SKU."""
    code = code.strip().upper()
    sql = f"""
    SELECT {SELECT_COLUMNS}
    FROM {TABLE_NAME}
    WHERE (UPPER(magento_ref_code) = %s
           OR UPPER(internal_ref_code) = %s
           OR UPPER(magento_ref_code) LIKE %s)
    ORDER BY quantity_sold DESC NULLS LAST
    LIMIT 20
    """
    return sql, [code, code, f"{code}%"]


# ======================================================
# Tool Entry Point
# ======================================================
@tool(args_schema=LeadSearchInput)
async def lead_search_tool(
    keywords: list[str] | None = None,
    product_line_vn: list[str] | None = None,
    gender_by_product: str | None = None,
    age_by_product: str | None = None,
    master_color: str | None = None,
    price_min: int | None = None,
    price_max: int | None = None,
    discount_min: int | None = None,
    discount_max: int | None = None,
    discovery_mode: str | None = None,
    size: str | None = None,
    magento_ref_code: str | None = None,
    reasoning: str | None = None,
) -> str:
    """
    Tim kiem san pham CANIFA - Hybrid: Keywords LIKE + Embedding cosine fallback.
    AI rut gon intent thanh keywords -> search description -> neu miss -> embed + cosine.
    """
    start = time.time()

    req = LeadSearchInput(
        keywords=keywords or [],
        product_line_vn=product_line_vn or [],
        gender_by_product=gender_by_product,
        age_by_product=age_by_product,
        master_color=master_color,
        price_min=price_min,
        price_max=price_max,
        discount_min=discount_min,
        discount_max=discount_max,
        discovery_mode=discovery_mode,
        size=size,
        magento_ref_code=magento_ref_code,
    )

    db = get_db_connection()

    try:
        fallback_msg = None
        s_details: dict = {}

        if req.magento_ref_code:
            sql, params = _build_sku_query(req.magento_ref_code)
            products = await db.execute_query_async(sql, params=tuple(params))
            tier = 0
        else:
            raw_query = " ".join(keywords) if keywords else (reasoning or "")
            products, tier, fallback_msg, s_details = await _cascading_search(req, db, raw_message=raw_query)

        elapsed_ms = round((time.time() - start) * 1000, 2)

        _search_modes = [
            "sku_lookup",
            "keywords_like",
            "embedding_cosine",
            "hard_filters_only",
            "drop_gender",
            "price_1.5x",
            "price_2x",
            "price_unlimited",
        ]
        _mode_label = _search_modes[min(tier, 7)]
        logger.info(
            "[LEAD SEARCH] ==========================================================",
        )
        logger.info(
            "[LEAD SEARCH] REQUEST  | keywords=%s | line=%s | gender=%s | color=%s | price=%s~%s | size=%s | mode=%s",
            keywords, product_line_vn, gender_by_product, master_color, price_min, price_max, size, discovery_mode,
        )
        logger.info(
            "[LEAD SEARCH] RESULT   | tier=%d (%s) | results=%d | time=%.0fms",
            tier, _mode_label, len(products), elapsed_ms,
        )
        if s_details.get("keyword_query"):
            logger.info("[LEAD SEARCH] KW-SEARCH | query=%s | results=%s | %.0fms",
                s_details["keyword_query"], s_details["keyword_results"], s_details.get("keyword_time_ms") or 0)
        if s_details.get("embed_query"):
            logger.info("[LEAD SEARCH] SEMANTIC  | query='%s' | results=%s | openai=%.0fms sr=%.0fms",
                s_details["embed_query"][:60], s_details["embed_results"],
                s_details.get("openai_embed_ms") or 0, s_details.get("starrocks_search_ms") or 0)
        logger.info("[LEAD SEARCH] ==========================================================")
        # Langfuse: log search result as a span on the existing trace
        if _LANGFUSE_AVAILABLE:
            try:
                _lf2 = get_langfuse()
                if _lf2:
                    with _lf2.start_as_current_observation(
                        as_type="span", name="lead_search_tool",
                        input={
                            "keywords": keywords or [],
                            "product_line_vn": product_line_vn or [],
                            "gender": gender_by_product,
                            "color": master_color,
                            "price_min": price_min,
                            "price_max": price_max,
                            "size": size,
                            "discovery_mode": discovery_mode,
                        },
                    ) as _root_sp:
                        _root_sp.update(
                            output={
                                "tier": tier,
                                "search_mode": _mode_label,
                                "results": len(products),
                                "elapsed_ms": elapsed_ms,
                            },
                            metadata={
                                "keyword_results": s_details.get("keyword_results"),
                                "keyword_time_ms": s_details.get("keyword_time_ms"),
                                "embed_query": s_details.get("embed_query"),
                                "embed_results": s_details.get("embed_results"),
                                "openai_embed_ms": s_details.get("openai_embed_ms"),
                                "starrocks_search_ms": s_details.get("starrocks_search_ms"),
                                "reasoning": reasoning,
                            },
                        )
            except Exception:
                pass

        formatted = await _format_products(products, db)

        result = {
            "status": "success",
            "count": len(products),
            "tier": tier,
            "search_mode": _search_modes[min(tier, 7)],
            "elapsed_ms": elapsed_ms,
            "reasoning": reasoning or "",
            "keywords_used": keywords or [],
            "products": formatted,
        }

        if req.magento_ref_code:
            result["search_details"] = {"mode": "sku_lookup"}
        else:
            result["search_details"] = {k: v for k, v in s_details.items() if v is not None}

        if fallback_msg:
            result["fallback_message"] = fallback_msg

        return json.dumps(result, ensure_ascii=False, default=str)

    except Exception as e:
        logger.error("Lead search tool error: %s", e, exc_info=True)
        return json.dumps({"status": "error", "message": str(e)})
