"""
Agent Helper Functions
Các hàm tiện ích cho chat controller.
"""

import json
import logging
import uuid
from decimal import Decimal

from langchain_core.messages import HumanMessage, ToolMessage
from langchain_core.runnables import RunnableConfig


from common.conversation_manager import ConversationManager
from common.langfuse_client import get_callback_handler
from common.starrocks_connection import get_db_connection

from .models import AgentState

import re

logger = logging.getLogger(__name__)


# ==============================================================================
# FORMAT PRODUCT RESULTS
# ==============================================================================


def _parse_description_text(desc: str) -> dict:
    """Parse description_text_full thành dict các field (backward compatibility)."""
    result = {}
    if not desc:
        return result

    name_match = re.search(r"product_name:\s*(.+?)\.(?:\s+master_color:|$)", desc)
    if name_match:
        result["product_name"] = name_match.group(1).strip()

    thumb_match = re.search(r"product_image_url_thumbnail:\s*(https?://[^\s]+?)\.(?:\s+product_web_url:|$)", desc)
    if thumb_match:
        result["product_image_url_thumbnail"] = thumb_match.group(1).strip()

    url_match = re.search(r"product_web_url:\s*(https?://[^\s]+?)\.(?:\s+description_text:|$)", desc)
    if url_match:
        result["product_web_url"] = url_match.group(1).strip()

    color_match = re.search(r"master_color:\s*(.+?)\.(?:\s+product_image_url:|$)", desc)
    if color_match:
        result["master_color"] = color_match.group(1).strip()

    return result


def format_product_results(products: list[dict]) -> list[dict]:
    """
    Format products - GROUP by base SKU (magento_ref_code), with multiple colors.

    Output format:
    {
        "sku": "1DS25S008",
        "name": "Váy liền bé gái",
        "colors": [
            {"color": "Cam/ Orange", "color_code": "SO123", "url": "...", "thumbnail": "..."},
            {"color": "Xanh/ Blue", "color_code": "SB456", "url": "...", "thumbnail": "..."}
        ],
        "price": 349000,
        "sale_price": 244000,
        "description": "..."
    }
    """
    max_products = 15
    grouped: dict[str, dict] = {}  # {magento_ref_code: {product_info}}

    for p in products:
        # Extract product info
        if p.get("product_name"):
            name = p["product_name"]
            color_name = p.get("master_color") or ""
            thumb_url = p.get("product_image_url_thumbnail") or ""
            web_url = p.get("product_web_url") or ""
            # Fallback: parse from description_text_full if fields are empty
            if not color_name or not thumb_url or not web_url:
                parsed = _parse_description_text(p.get("description_text_full", ""))
                color_name = color_name or parsed.get("master_color", "")
                thumb_url = thumb_url or parsed.get("product_image_url_thumbnail", "")
                web_url = web_url or parsed.get("product_web_url", "")
        else:
            desc_full = p.get("description_text_full", "")
            parsed = _parse_description_text(desc_full)
            name = parsed.get("product_name", "")
            color_name = parsed.get("master_color", "")
            thumb_url = parsed.get("product_image_url_thumbnail", "")
            web_url = parsed.get("product_web_url", "")

        original_price = p.get("original_price") or 0
        sale_price = p.get("sale_price") or 0
        magento_ref = p.get("magento_ref_code", "")
        product_color_code = p.get("product_color_code", "")

        # Extract color code from product_color_code (VD: 1DS25S008-SO123 → SO123)
        color_code_only = ""
        if product_color_code and "-" in product_color_code:
            parts = product_color_code.split("-", 1)
            color_code_only = parts[1] if len(parts) > 1 else ""

        # Use magento_ref as base SKU for grouping
        base_sku = magento_ref if magento_ref else product_color_code
        if not base_sku:
            continue

        # Color variant info
        color_variant = {
            "color": color_name,
            "color_code": color_code_only,
            "url": web_url,
            "thumbnail": thumb_url,
        }

        if base_sku in grouped:
            # Add color to existing product
            existing_colors = [c["color"] for c in grouped[base_sku]["colors"]]
            if color_name and color_name not in existing_colors:
                grouped[base_sku]["colors"].append(color_variant)
                # Update price range if different
                if sale_price and sale_price < grouped[base_sku].get("sale_price", float("inf")):
                    grouped[base_sku]["sale_price"] = int(sale_price)
        else:
            # New product - use first color's URL/thumbnail as default
            product_entry = {
                "sku": base_sku,
                "name": name,
                "color": color_name,  # First color as default
                "colors": [color_variant] if color_name else [],
                "price": int(original_price),
                "sale_price": int(sale_price) if sale_price else int(original_price),
                "url": web_url,  # First color's URL
                "thumbnail_image_url": thumb_url,  # First color's thumbnail
                "description": (p.get("description_text") or "")[:200],
            }
            # Include sizes if available (pipe-separated → list)
            size_scale = p.get("size_scale")
            if size_scale:
                product_entry["sizes"] = [s.strip() for s in size_scale.split("|") if s.strip()]
            # Include quantity_sold if available (for best seller)
            qty_sold = p.get("quantity_sold")
            if qty_sold is not None:
                product_entry["quantity_sold"] = int(qty_sold)
            grouped[base_sku] = product_entry

    formatted = list(grouped.values())[:max_products]
    logger.info(f"📦 Formatted {len(formatted)} products (grouped by SKU)")
    return formatted



def decimal_default(obj):
    """
    JSON serializer for objects not serializable by default json code.
    Handles Decimal objects.
    """
    if isinstance(obj, Decimal):
        return float(obj)
    raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")


def extract_product_ids(messages: list) -> list[dict]:
    """
    Extract full product info from tool messages (data_retrieval_tool results).
    Returns list of product objects with: sku, name, price, sale_price, url, thumbnail_image_url.
    """
    products = []
    seen_skus = set()

    for msg in messages:
        if isinstance(msg, ToolMessage):
            try:
                # Tool result is JSON string
                tool_result = json.loads(msg.content)

                # Check if tool returned products (new format with "results" wrapper)
                if tool_result.get("status") == "success":
                    # Handle both direct "products" and nested "results" format
                    product_list = []

                    if "results" in tool_result:
                        results_data = tool_result["results"]
                        if results_data and isinstance(results_data, list):
                            # Check first item to determine format
                            first_item = results_data[0] if len(results_data) > 0 else {}
                            if isinstance(first_item, dict) and "products" in first_item:
                                # Nested format: {"results": [{"products": [...]}]}
                                for result_item in results_data:
                                    product_list.extend(result_item.get("products", []))
                            else:
                                # Flat format: {"results": [product1, product2]} (Current)
                                product_list = results_data
                    elif "products" in tool_result:
                        # Legacy format: {"products": [...]}
                        product_list = tool_result["products"]

                    logger.warning(f"🛠️ [EXTRACT] Extracted {len(product_list)} products from tool")

                    for product in product_list:
                        # ⚡ FLATTEN: Tách variants thành separate products (mỗi màu = 1 product)
                        # Grouped: {"product_id": "...", "name": "...", "variants": [...]}
                        # → Tách thành N products (1 per variant)

                        # Check if grouped format (has variants)
                        if "variants" in product and product.get("variants"):
                            # Grouped product - expand EACH variant into separate product
                            product_name = product.get("name", "")
                            base_sku = product.get("product_id")  # ✅ Giữ base SKU giống nhau

                            for variant in product["variants"]:
                                variant_sku = variant.get("sku")
                                # ✅ Use base_sku instead of variant_sku for consistency
                                display_sku = base_sku if base_sku else variant_sku

                                # Create unique key for dedup using variant SKU
                                dedup_key = variant_sku or display_sku

                                if dedup_key and dedup_key not in seen_skus:
                                    seen_skus.add(dedup_key)

                                    # ✅ Spread ALL variant fields first, then override sku + name
                                    product_obj = {
                                        **variant,  # Copy all variant fields (color, price, discount, stock, url, thumbnail, etc.)
                                        "sku": display_sku,  # Override with base SKU
                                        "name": product_name,  # Override with product name
                                    }
                                    products.append(product_obj)
                        else:
                            # Flat format - use directly
                            sku = product.get("sku") or product.get("internal_ref_code")

                            if sku and sku not in seen_skus:
                                seen_skus.add(sku)

                                product_obj = {
                                    "sku": sku,
                                    "name": product.get("name", ""),
                                    "color": product.get("color", ""),
                                    "price": product.get("price", 0),
                                    "sale_price": product.get("sale_price"),
                                    "url": product.get("url", ""),
                                    "thumbnail_image_url": product.get("thumbnail_image_url", ""),
                                }
                                products.append(product_obj)
            except (json.JSONDecodeError, KeyError, TypeError) as e:
                logger.debug(f"Could not parse tool message for products: {e}")
                continue

    return products


def parse_ai_response_fast(ai_raw_content: str) -> tuple[str, list[str], str | None]:
    """
    FAST parse - Chỉ extract ai_response + product_ids từ JSON, KHÔNG query DB.
    Trả về SKU list thay vì full product objects.

    Returns:
        tuple: (ai_text_response, product_skus, user_insight_json)
    """
    import json
    import re

    ai_text_response = ai_raw_content
    product_skus = []
    user_insight = None

    try:
        ai_json = json.loads(ai_raw_content)

        # Extract basic fields
        ai_text_response = ai_json.get("ai_response", ai_raw_content)
        explicit_skus = ai_json.get("product_ids", [])
        raw_insight = ai_json.get("user_insight")

        # Extract SKUs mentioned in text
        mentioned_skus_in_text = set(re.findall(r"\[([A-Z0-9]+)\]", ai_text_response))

        # Determine target SKUs
        if explicit_skus and isinstance(explicit_skus, list):
            product_skus = [str(s) for s in explicit_skus]
        elif mentioned_skus_in_text:
            product_skus = list(mentioned_skus_in_text)

        # Convert user_insight to JSON string
        if raw_insight:
            if isinstance(raw_insight, dict):
                user_insight = json.dumps(raw_insight, ensure_ascii=False, indent=2)
            elif isinstance(raw_insight, str):
                user_insight = raw_insight

        logger.info(f"⚡ Fast parse: ai_response={len(ai_text_response)} chars, skus={product_skus}")

    except (json.JSONDecodeError, TypeError) as e:
        logger.warning(f"⚠️ Fast parse failed: {e}")

    return ai_text_response, product_skus, user_insight


async def parse_ai_response_async(ai_raw_content: str, all_products: list) -> tuple[str, list, str | None]:
    """
    Async version of parse_ai_response with DB fallback.

    Parse AI response từ LLM output và map SKUs với product data.
    Nếu SKU được mention nhưng không có trong all_products (context hiện tại),
    sẽ query trực tiếp DB để lấy thông tin.

    Flow:
    - LLM trả về: {"ai_response": "...", "product_ids": ["SKU1"], ...}
    - Map SKUs → enriched products từ context
    - Nếu thiếu → Query DB
    """
    import re

    from .structured_models import ChatResponse

    ai_text_response = ai_raw_content
    final_products = []
    user_insight = None

    logger.info(f"🤖 Raw AI JSON: {ai_raw_content}")

    try:
        # Try to parse if it's a JSON string from LLM
        ai_json = json.loads(ai_raw_content)

        # === PYDANTIC VALIDATION ===
        try:
            # Try strict Pydantic validation
            parsed_response = ChatResponse.model_validate(ai_json)
            ai_text_response = parsed_response.ai_response
            explicit_skus = parsed_response.product_ids

            # Convert user_insight to dict/string for storage
            if parsed_response.user_insight:
                user_insight = parsed_response.user_insight.model_dump_json(indent=2)

            logger.info("✅ Pydantic validation passed for ChatResponse")

        except Exception as validation_error:
            # Fallback to manual parsing if Pydantic fails
            logger.warning(f"⚠️ Pydantic validation failed, using fallback: {validation_error}")
            ai_text_response = ai_json.get("ai_response", ai_raw_content)
            explicit_skus = ai_json.get("product_ids", [])
            raw_insight = ai_json.get("user_insight")

            if raw_insight:
                if isinstance(raw_insight, dict):
                    user_insight = json.dumps(raw_insight, ensure_ascii=False, indent=2)
                elif isinstance(raw_insight, str):
                    user_insight = raw_insight

        # === CRITICAL: Filter/Fetch products ===
        # Extract SKUs mentioned in ai_response text using regex pattern [SKU]
        mentioned_skus_in_text = set(re.findall(r"\[([A-Z0-9]+)\]", ai_text_response))
        logger.info(f"📝 SKUs mentioned in ai_response: {mentioned_skus_in_text}")

        # Determine target SKUs
        target_skus = set()

        # 1. Use explicit SKUs if available and confirmed by text, OR just explicit
        if explicit_skus and isinstance(explicit_skus, list):
            # Optional: Filter explicit SKUs to only those actually in text to reduce hallucination
            # But if explicit list is provided, we generally trust it unless we want strict text-match
            if mentioned_skus_in_text:
                explicit_set = set(str(s) for s in explicit_skus)
                target_skus = explicit_set.intersection(mentioned_skus_in_text)
                if not target_skus:  # If intersection empty, fallback to text mentions
                    target_skus = mentioned_skus_in_text
            else:
                target_skus = set(str(s) for s in explicit_skus)
        elif mentioned_skus_in_text:
            # 2. If no explicit SKUs, use text mentions
            target_skus = mentioned_skus_in_text

        logger.info(f"🎯 Target SKUs to return: {target_skus}")

        if target_skus:
            # Build lookup from current context
            product_lookup = {p["sku"]: p for p in all_products if p.get("sku")}

            found_products = []


            for sku in target_skus:
                if sku in product_lookup:
                    found_products.append(product_lookup[sku])
                else:
                    # SKU not in context (e.g., only stock check was called)
                    # Don't create dummy product - just skip
                    logger.debug(f"⚠️ SKU {sku} not in context (stock-only query?), skipping")

            final_products = found_products

    except (json.JSONDecodeError, TypeError) as e:
        logger.warning(f"⚠️ Failed to parse AI response as JSON: {e}")

    return ai_text_response, final_products, user_insight


def prepare_execution_context(query: str, user_id: str, history: list, images: list | None):
    """
    Prepare initial state and execution config for the graph run.

    Returns:
        tuple: (initial_state, exec_config)
    """
    initial_state: AgentState = {
        "user_query": HumanMessage(content=query),
        "messages": [HumanMessage(content=query)],
        "history": history,
        "user_id": user_id,
        "images_embedding": [],
        "ai_response": None,
    }
    run_id = str(uuid.uuid4())

    # Metadata for LangChain (tags for logging/filtering)
    metadata = {
        "run_id": run_id,
        "tags": "chatbot,production",
    }

    langfuse_handler = get_callback_handler()

    exec_config = RunnableConfig(
        configurable={
            "user_id": user_id,
            "transient_images": images or [],
            "run_id": run_id,
        },
        run_id=run_id,
        metadata=metadata,
        callbacks=[langfuse_handler] if langfuse_handler else [],
    )
    return initial_state, exec_config


async def handle_post_chat_async(
    memory: ConversationManager, identity_key: str, human_query: str, ai_response: dict | None, conversation_id: str | None = None
):
    """
    Save chat history in background task after response is sent.
    Lưu AI response dưới dạng JSON string.
    """
    if ai_response:
        try:
            # Convert dict thành JSON string để lưu vào TEXT field
            # Use decimal_default to handle Decimal types from DB
            ai_response_json = json.dumps(ai_response, ensure_ascii=False, default=decimal_default)
            await memory.save_conversation_turn(identity_key, human_query, ai_response_json, conversation_id=conversation_id)
            logger.debug(f"Saved conversation for identity_key {identity_key} (conv: {conversation_id})")
        except Exception as e:
            logger.error(f"Failed to save conversation for identity_key {identity_key}: {e}", exc_info=True)
