# src/layout_generator/assets.py

"""
Шаг пайплайна для загрузки и кэширования 3D-моделей инвентаря.
"""
import logging
from typing import Any, Dict, List

from dsynth.assets.asset import load_assets_lib
from dsynth.scene_gen.utils import flatten_dict

from .base import BaseStep, LayoutContext

logger: logging.Logger = logging.getLogger(__name__)

# Глобальный кэш ОЗУ
_ASSET_CACHE: Dict[str, Any] = {}

class LoadAssetsStep(BaseStep):
    """Шаг предзагрузки 3D-моделей в оперативную память."""
    
    def process(self, context: LayoutContext) -> LayoutContext:
        global _ASSET_CACHE
        if _ASSET_CACHE:
            logger.info("⚡ Используем модели из кэша ОЗУ.")
            context.product_assets_lib = _ASSET_CACHE['products']
            context.large_items = _ASSET_CACHE['large']
            context.medium_items = _ASSET_CACHE['medium']
            context.small_items = _ASSET_CACHE['small']
            return context
            
        logger.info("⏳ ТОТАЛЬНЫЙ РАЗОГРЕВ КЭША: Начинаем чтение геометрии с диска...")
        
        short_assets_lib: Dict[str, Any] = flatten_dict(load_assets_lib(context.cfg.assets.products_hierarchy), sep='.')
        valid_food_keys: List[str] = [k for k in short_assets_lib.keys() if "food" in k.lower()]

        if not valid_food_keys:
            logger.critical("В библиотеке отсутствуют модели еды!")
            return context

        large_kw: List[str] = ['2l', '1.5l', '1l', 'bottle', 'box', 'milk', 'cereal', 'detergent', 'wine', 'vodka', 'juice', 'shampoo', 'lotion']
        small_kw: List[str] = ['gum', 'can', 'yogurt', 'tuna', 'spice', 'bar', 'candy', 'soup', 'noodle', 'butter', 'salt']

        large_items: List[str] = []
        small_items: List[str] = []
        medium_items: List[str] = []

        for k in valid_food_keys:
            k_lower: str = k.lower()
            if any(w in k_lower for w in large_kw):
                large_items.append(k)
            elif any(w in k_lower for w in small_kw):
                small_items.append(k)
            else:
                medium_items.append(k)

        product_assets_lib: Dict[str, Any] = {}
        for k, v in short_assets_lib.items():
            product_assets_lib[k] = v
            product_assets_lib[f"products_hierarchy.{k}"] = v

        if "fixtures" in context.cfg.assets:
            fixtures_lib: Dict[str, Any] = flatten_dict(load_assets_lib(context.cfg.assets.fixtures), sep='.')
            for k, v in fixtures_lib.items():
                product_assets_lib[f"fixtures.{k}"] = v
                product_assets_lib[k] = v

        total_products: int = len(valid_food_keys)
        logger.info(f"Найдено {total_products} моделей инвентаря. Приступаем к парсингу...")
        
        loaded_count: int = 0
        for item_key in valid_food_keys:
            try:
                asset_obj: Any = product_assets_lib.get(item_key)
                if asset_obj:
                    _ = asset_obj.ss_asset
                    loaded_count += 1
                    if loaded_count % 100 == 0 or loaded_count == total_products:
                        logger.info(f"   📦 Прогресс: Загружено {loaded_count} из {total_products} товаров...")
            except Exception as e:
                logger.warning(f"⚠️ Ошибка чтения 3D-модели {item_key}: {e}")
                
        logger.info("Начинаем загрузку 3D-моделей стеллажей и оборудования...")
        loaded_fixtures: int = 0
        for shelf in ["fixtures.large_showcase_fake", "fixtures.small_shelf_two_sided", "fixtures.shelf_metal", "fixtures.shelf_low_res"]:
            try:
                asset_obj = product_assets_lib.get(shelf)
                if asset_obj:
                    _ = asset_obj.ss_asset
                    loaded_fixtures += 1
            except Exception:
                pass
                
        logger.info(f"   🗄️ Прогресс: Загружено {loaded_fixtures} моделей оборудования.")

        _ASSET_CACHE['products'] = product_assets_lib
        _ASSET_CACHE['large'] = large_items
        _ASSET_CACHE['medium'] = medium_items
        _ASSET_CACHE['small'] = small_items
        
        logger.info(f"✅ Тотальный разогрев успешно завершен! В ОЗУ надежно зафиксировано {loaded_count + loaded_fixtures} 3D-моделей.")
        
        context.product_assets_lib = product_assets_lib
        context.large_items = large_items
        context.medium_items = medium_items
        context.small_items = small_items
        
        return context