
import json

from langchain_core.messages import ToolMessage

from agent.helper import extract_product_ids

# Mock ToolMessage from sku_search_tool (Current buggy format)
sku_search_output = {
    "status": "success",
    "products": [
        {
            "sku": "6TS25S003",
            "sku_color": "6TS25S003-SB055",
            "name": "Áo thun nam",
            "image": "https://canifa.com/image1.jpg",
            "url": "https://canifa.com/p1"
        }
    ]
}

# Mock ToolMessage from data_retrieval_tool (Working format)
data_retrieval_output = {
    "status": "success",
    "results": [
        {
            "sku": "6TS25S003",
            "name": "Áo thun nam",
            "thumbnail_image_url": "https://canifa.com/image1.jpg",
            "url": "https://canifa.com/p1"
        }
    ]
}

def test_extraction():
    print("Testing SKU Search Tool Extraction...")
    msg_sku = ToolMessage(content=json.dumps(sku_search_output), tool_call_id="1")
    products_sku = extract_product_ids([msg_sku])
    print(f"Extracted from SKU Search: {json.dumps(products_sku, indent=2, ensure_ascii=False)}")
    
    # Check if thumbnail_image_url is present
    if products_sku and products_sku[0].get("thumbnail_image_url"):
        print("✅ SKU Search Extraction: SUCCESS")
    else:
        print("❌ SKU Search Extraction: FAILED (thumbnail_image_url missing)")

    print("\nTesting Data Retrieval Tool Extraction...")
    msg_data = ToolMessage(content=json.dumps(data_retrieval_output), tool_call_id="2")
    products_data = extract_product_ids([msg_data])
    print(f"Extracted from Data Retrieval: {json.dumps(products_data, indent=2, ensure_ascii=False)}")
    
    if products_data and products_data[0].get("thumbnail_image_url"):
        print("✅ Data Retrieval Extraction: SUCCESS")
    else:
        print("❌ Data Retrieval Extraction: FAILED")

if __name__ == "__main__":
    test_extraction()
