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

Tracing:
- LangChain CallbackHandler auto-traces tool call (name, input, output)
- Manual span 'db-query-search' in _execute_single_search for SQL timing & product count
  (nested under parent LangChain trace via start_as_current_observation)
"""

import asyncio
import json
import logging
import time


from langfuse import get_client as get_langfuse

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


from agent.helper import format_product_results
from agent.tools.product_search_helpers import build_starrocks_query
from agent.tools.stock_helpers import fetch_stock_for_skus
from common.starrocks_connection import get_db_connection

# Setup Logger
logger = logging.getLogger(__name__)




class SearchItem(BaseModel):
    model_config = {"extra": "ignore"}  # Gemini may send extra fields

    # ====== SEMANTIC SEARCH ======
    description: str = Field(
        description=(
            "Mô tả sản phẩm cho semantic search. "
            "FORMAT: 'product_name: [tên]. description_text: [mô tả]. "
            "material_group: [chất liệu]. style: [phong cách]. fitting: [dáng]. "
            "form_sleeve: [tay]. form_neckline: [cổ]. product_line_vn: [dòng SP]'. "
            "🚨 product_name = CHÍNH XÁC lời user nói, TUYỆT ĐỐI KHÔNG tự đổi sang tên khác. "
            "⚠️ KHÔNG đưa master_color vào description — dùng field master_color riêng! "
        )
    )
    product_name: str | None = Field(
        default=None,
        description="Tên sản phẩm. Chỉ điền khi khách cung cấp tên cụ thể.",
    )

    # ====== SKU LOOKUP ======
    magento_ref_code: str | None = Field(
        default=None,
        description="Mã sản phẩm chính xác (SKU). Chỉ điền khi khách cung cấp mã cụ thể.",
    )

    # ====== SQL HARD FILTERS ======
    gender_by_product: str | None = Field(
        default=None,
        description="[SQL FILTER] Giới tính. GIÁ TRỊ HỢP LỆ: women, men, boy, girl, unisex, newborn",
    )
    age_by_product: str | None = Field(
        default=None,
        description="[SQL FILTER] Độ tuổi. GIÁ TRỊ HỢP LỆ: adult, kid, others",
    )
    master_color: str | None = Field(
        default=None,
        description=(
            "[SQL FILTER] Màu sắc sản phẩm. Gửi CHÍNH XÁC từ khách nói. "
            "Tool sẽ tự match trong DB bằng LIKE. KHÔNG đưa màu vào description!"
        ),
    )
    price_min: int | None = Field(
        default=None,
        description="[SQL FILTER] Giá thấp nhất (VND)",
    )
    price_max: int | None = Field(
        default=None,
        description="[SQL FILTER] Giá cao nhất (VND)",
    )
    discount_min: int | None = Field(
        default=None,
        description="[SQL FILTER] % giảm giá tối thiểu. VD: 50 → chỉ lấy SP giảm >= 50%.",
    )
    discount_max: int | None = Field(
        default=None,
        description="[SQL FILTER] % giảm giá tối đa. VD: 70 → chỉ lấy SP giảm <= 70%.",
    )
    discovery_mode: str | None = Field(
        default=None,
        description=(
            "[SQL FILTER] 'new' = hàng mới, 'best_seller' = bán chạy. "
            "Chỉ dùng khi khách NÓI RÕ. Nếu không nói → null."
        ),
    )

    @classmethod
    def __get_validators__(cls):
        yield cls._coerce_empty_strings

    def model_post_init(self, __context) -> None:
        """Coerce empty strings to None and extract product_name from description."""
        import re
        # 1. Empty string → None
        for field_name in ['product_name', 'magento_ref_code', 'gender_by_product',
                           'age_by_product', 'master_color', 'discovery_mode']:
            val = getattr(self, field_name, None)
            if isinstance(val, str) and val.strip() == '':
                object.__setattr__(self, field_name, None)

        # 2. If product_name is None but description contains "product_name: X" → extract it
        if self.product_name is None and self.description:
            match = re.search(r'product_name:\s*([^.]+?)(?:\.|$)', self.description)
            if match:
                extracted = match.group(1).strip()
                if extracted and extracted.lower() not in ('sản phẩm', 'product'):
                    object.__setattr__(self, 'product_name', extracted)
                    logger.info(f"📝 Extracted product_name from description: '{extracted}'")


class MultiSearchParams(BaseModel):
    model_config = {"extra": "ignore"}  # Gemini may send extra fields
    searches: list[SearchItem] = Field(description="Danh sách các truy vấn tìm kiếm")


async def _execute_single_search(
    db, item: SearchItem, query_vector: list[float] | None = None
) -> tuple[list[dict], dict]:
    """
    Thực thi một search query đơn lẻ (Async).
    Creates a manual 'db-query-search' span under the current Langfuse context.

    Returns:
        Tuple of (products, filter_info)
    """
    try:
        short_query = (
            (item.description[:60] + "...") if item.description and len(item.description) > 60 else item.description
        )
        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, params = 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)

        if not sql:
            return [], {"fallback_used": False}

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

        # Debug: Log first product to see fields
        if products:
            first_p = products[0]
            logger.debug("🔍 [DEBUG] First product keys: %s", list(first_p.keys()))
            logger.debug(
                "🔍 [DEBUG] First product price: %s, sale_price: %s",
                first_p.get("original_price"),
                first_p.get("sale_price"),
            )
            logger.debug(
                "🔍 [DEBUG] size_scale: %r, description_text: %r",
                first_p.get("size_scale"),
                (first_p.get("description_text") or "")[:100],
            )

        # ✅ Nest 'db-query-search' span dưới chat-request trace hiện tại trên Langfuse
        try:
            lf = get_langfuse()
            if lf:
                with lf.start_as_current_observation(
                    as_type="span",
                    name="db-query-search",
                ) as span:
                    span.update(
                        input={
                            "description": item.description[:200] if item.description else None,
                            "product_name": item.product_name,
                            "gender": item.gender_by_product,
                            "age": item.age_by_product,
                            "color": item.master_color,
                            "price_range": f"{item.price_min or ''}-{item.price_max or ''}",
                            "magento_ref_code": item.magento_ref_code,
                            "discovery_mode": item.discovery_mode,
                        },
                        output={
                            "product_count": len(products),
                            "build_ms": round(query_build_time, 2),
                            "db_ms": round(db_time, 2),
                            "total_ms": round(query_build_time + db_time, 2),
                            "sample_products": [
                                {
                                    "name": p.get("product_name", "?"),
                                    "code": p.get("magento_ref_code", "?"),
                                    "price": float(p["sale_price"]) if p.get("sale_price") else None,
                                }
                                for p in products[:3]
                            ] if products else [],
                        },
                        metadata={"sql_length": len(sql)},
                    )
        except Exception as trace_err:
            logger.warning("⚠️ Langfuse db-query-search span failed: %s", trace_err, exc_info=True)

        formatted = format_product_results(products)

        # ═══ AUTO STOCK CHECK: Gọi Stock API rồi gắn is_in_stock vào product ═══
        try:
            # Collect product_color_codes từ raw products (trước khi group)
            color_codes = []
            for p in products:
                pcc = p.get("product_color_code", "")
                if pcc:
                    color_codes.append(pcc)

            if color_codes:
                stock_map = await fetch_stock_for_skus(color_codes, timeout_sec=3.0)

                for product in formatted:
                    # Check tất cả color variants của product
                    any_in_stock = False
                    stock_sizes = []

                    for color_variant in product.get("colors", []):
                        # Build full product_color_code: SKU-COLOR_CODE
                        full_code = f"{product['sku']}-{color_variant.get('color_code', '')}"

                        # Check all sizes for this color
                        for sku_key, stock_info in stock_map.items():
                            if sku_key.startswith(full_code):
                                if stock_info.get("is_in_stock"):
                                    any_in_stock = True
                                    # Extract size from SKU (e.g. 6TS25S018-SY322-L → L)
                                    size_part = sku_key.replace(full_code + "-", "") if "-" in sku_key[len(full_code):] else ""
                                    if size_part and stock_info.get("qty", 0) > 0:
                                        stock_sizes.append(size_part)

                    product["is_in_stock"] = any_in_stock
                    if stock_sizes:
                        product["available_sizes"] = stock_sizes

                logger.info("📦 [STOCK] Enriched %d products with stock info", len(formatted))
        except Exception as stock_err:
            logger.warning("⚠️ Stock enrichment failed (non-fatal): %s", stock_err)
            # Không crash — products vẫn trả về, chỉ thiếu stock info

        return formatted, {"fallback_used": False}
    except Exception as e:
        logger.exception("Single search error for item %r: %s", item, e)
        return [], {"fallback_used": False, "error": str(e)}


@tool(args_schema=MultiSearchParams)
async def data_retrieval_tool(searches: list[SearchItem]) -> str:
    """
    Công cụ tìm kiếm sản phẩm CANIFA.
    Hỗ trợ tìm kiếm Semantic và lọc theo Metadata.
    """
    tool_start = time.time()
    logger.info("🔧 data_retrieval_tool called with %d items", len(searches))

    # Log search params cho debugging
    for i, s in enumerate(searches):
        logger.info(
            "🔧 Search[%d]: desc=%r, name=%r, gender=%s, age=%s, color=%s, "
            "price=%s-%s, code=%s, mode=%s",
            i,
            (s.description[:80] + "...") if s.description and len(s.description) > 80 else s.description,
            s.product_name,
            s.gender_by_product,
            s.age_by_product,
            s.master_color,
            s.price_min, s.price_max,
            s.magento_ref_code,
            s.discovery_mode,
        )

    # Get DB Connection
    db = get_db_connection()
    if not db:
        return json.dumps({"status": "error", "message": "Database connection failed"})

    combined_results = []
    all_filter_infos = []

    tasks = []
    for item in searches:
        tasks.append(_execute_single_search(db, item))

    results_list = await asyncio.gather(*tasks)

    for products, filter_info in results_list:
        combined_results.extend(products)
        if filter_info:
            all_filter_infos.append(filter_info)

    # Aggregate filter info from first result for simplicity in response
    final_info = all_filter_infos[0] if all_filter_infos else {}

    # Include search input so LLM knows its own reasoning
    search_inputs = [
        {
            "description": item.description,
            "gender_by_product": item.gender_by_product,
            "age_by_product": item.age_by_product,
            "discovery_mode": item.discovery_mode,
        }
        for item in searches
    ]

    output = {
        "status": "success",
        "search_input": search_inputs,
        "results": combined_results,
        "filter_info": final_info,
    }

    total_ms = (time.time() - tool_start) * 1000

    logger.info(
        "🎁 Final result: %d products, total_ms=%.2f. Fallback used: %s",
        len(combined_results),
        total_ms,
        final_info.get("fallback_used", False),
    )

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


