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

import json
import logging
import uuid

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

from common.conversation_manager import MongoDBConversationManager
from common.langfuse_client import get_callback_handler
from .models import AgentState

logger = logging.getLogger(__name__)


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
                if tool_result.get("status") == "success" and "products" in tool_result:
                    for product in tool_result["products"]:
                        sku = product.get("internal_ref_code")
                        if sku and sku not in seen_skus:
                            seen_skus.add(sku)

                            # Extract full product info
                            product_obj = {
                                "sku": sku,
                                "name": product.get("magento_product_name", ""),
                                "price": product.get("price_vnd", 0),
                                "sale_price": product.get("sale_price_vnd"),  # null nếu không sale
                                "url": product.get("magento_url_key", ""),
                                "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(ai_raw_content: str, all_product_ids: list) -> tuple[str, list]:
    """
    Parse AI response từ LLM output.
    
    Args:
        ai_raw_content: Raw content từ AI response
        all_product_ids: Product IDs extracted từ tool messages
        
    Returns:
        tuple: (ai_text_response, final_product_ids)
    """
    ai_text_response = ai_raw_content
    final_product_ids = all_product_ids

    try:
        # Strip markdown code blocks if present (```json ... ```)
        content_to_parse = ai_raw_content.strip()
        
        # Remove markdown code block markers
        if content_to_parse.startswith("```"):
            # Find the first newline after ```
            first_newline = content_to_parse.find("\n")
            if first_newline != -1:
                # Remove opening ```json or ``` and closing ```
                content_to_parse = content_to_parse[first_newline + 1:]
                if content_to_parse.endswith("```"):
                    content_to_parse = content_to_parse[:-3]
                content_to_parse = content_to_parse.strip()
        
        # Try to parse if it's a JSON string from LLM
        ai_json = json.loads(content_to_parse)
        ai_text_response = ai_json.get("ai_response", ai_raw_content)
        explicit_ids = ai_json.get("product_ids", [])
        if explicit_ids and isinstance(explicit_ids, list):
            # Replace with explicit IDs from LLM
            final_product_ids = explicit_ids
    except (json.JSONDecodeError, TypeError, ValueError):
        # If parsing fails, return original content
        pass

    return ai_text_response, final_product_ids


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: MongoDBConversationManager, 
    identity_key: str, 
    human_query: str, 
    ai_response: dict | 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
            ai_response_json = json.dumps(ai_response, ensure_ascii=False)
            await memory.save_conversation_turn(identity_key, human_query, ai_response_json)
            logger.debug(f"Saved conversation for identity_key {identity_key}")
        except Exception as e:
            logger.error(f"Failed to save conversation for identity_key {identity_key}: {e}", exc_info=True)
