import asyncio
import json
import logging

from agent.tools.check_is_stock import check_is_stock
from agent.tools.data_retrieval_tool import data_retrieval_tool
from common.starrocks_connection import StarRocksConnection

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def test_full_flow():
    print("🚀 STARTING STOCK FLOW VERIFICATION\n")
    
    # CASE 0: Verify Prompt Loading
    print("--- TEST CASE 0: Tool Description Loading ---")
    if "KIỂM TRA TỒN KHO" in check_is_stock.description:
        print("✅ check_is_stock description updated correctly.")
    else:
        print(f"❌ check_is_stock description mismatch. Got: {check_is_stock.description[:50]}...")
        
    if "Siêu công cụ tìm kiếm" in data_retrieval_tool.description:
        print("✅ data_retrieval_tool description updated correctly.")
    else:
        print(f"❌ data_retrieval_tool description mismatch. Got: {data_retrieval_tool.description[:50]}...")

    # CASE 1: Specific SKU Check (check_is_stock)
    sku_input = "6ST25W005-SE091" 
    print(f"\n--- TEST CASE 1: Check Specific SKU: {sku_input} ---")
    try:
        # Tool call via ainvoke
        res = await check_is_stock.ainvoke({"skus": sku_input})
        res_json = json.loads(res)
        if "stock_responses" in res_json:
            print(f"✅ Tool returned stock_responses. Count: {len(res_json['stock_responses'])}")
            print(f"   Sample: {res_json['stock_responses'][:1]}")
        else:
            print(f"❌ Tool response missing stock_responses: {res}")
    except Exception as e:
         print(f"❌ Error calling check_is_stock: {e}")

    # CASE 2: Base Code Expansion (check_is_stock)
    base_code = "6TE25W005"
    print(f"\n--- TEST CASE 2: Check Base Code Expansion: {base_code} ---")
    try:
        res = await check_is_stock.ainvoke({"skus": base_code})
        res_json = json.loads(res)
        if "stock_responses" in res_json and len(res_json['stock_responses']) > 1:
             print(f"✅ Tool expanded base code. Count: {len(res_json['stock_responses'])}")
        else:
             print(f"⚠️ Tool might not have expanded or code invalid. Count: {len(res_json.get('stock_responses', []))}")
             print(f"   Result: {res_json}")
    except Exception as e:
        print(f"❌ Error calling check_is_stock with base code: {e}")

    # CASE 3: Data Retrieval Integration
    query_text = "áo thun nam"
    print(f"\n--- TEST CASE 3: Data Retrieval + Stock Enrichment ({query_text}) ---")
    try:
        # Mocking a search item as DICT with ALL keys (Strict Mode)
        search_item_dict = {
            "description": "Áo thun nam",
            "product_name": "Áo thun",
            "magento_ref_code": None,
            "price_min": None,
            "price_max": None,
            "gender_by_product": None,
            "age_by_product": None,
            "master_color": None,
            "form_sleeve": None,
            "style": None,
            "fitting": None,
            "form_neckline": None,
            "material_group": None,
            "season": None,
            "product_line_vn": None
        }
        # Tool call via ainvoke with DICT input
        res_str = await data_retrieval_tool.ainvoke({"searches": [search_item_dict]})
        res_json = json.loads(res_str)
        
        results = res_json.get("results", [])
        print(f"Found {len(results)} products.")
        if results:
            print(f"Sample Product Keys: {results[0].keys()}")
            print(f"Sample Product Code: {results[0].get('product_color_code')}")
            print(f"Sample Magento Code: {results[0].get('magento_ref_code')}")
        
        enriched_count = 0
        for p in results[:5]: # Check first 5
            stock = p.get("stock_info")
            if stock:
                enriched_count += 1
                # print(f"   Product {p.get('product_color_code')}: Has Stock Info ✅")
            else:
                pass
                # print(f"   Product {p.get('product_color_code')}: No Stock Info ❌")
        
        if res_json.get("stock_enriched") is True:
            print(f"✅ Response has 'stock_enriched'=True")
        else:
            print(f"❌ Response 'stock_enriched' is False/Missing")
            
        print(f"✅ Verified stock info present in {enriched_count}/5 sample items.")
        
    except Exception as e:
        print(f"❌ Error calling data_retrieval_tool: {e}")
        if hasattr(e, 'errors'):
             print(f"   Details: {e.errors()}")

    # Cleanup
    await StarRocksConnection.clear_pool()
    print("\n🚀 VERIFICATION COMPLETE")

if __name__ == "__main__":
    asyncio.run(test_full_flow())
