src/ui/callbacks.py

# src/ui/callbacks.py

"""
Модуль обработчиков событий (Controller) для веб-интерфейса Gradio.
Отвечает за связывание пользовательского ввода с LLM-клиентом и 3D-пайплайном.
"""

import logging
import queue
import threading
import time
from pathlib import Path
from typing import Any, Dict, Generator, List, Optional, Tuple

import gradio as gr
from omegaconf import DictConfig, OmegaConf

from layout_generator.pipeline import generate_layout_stream
from llm.client import InventoryLLMClient
from utils.logger import setup_logger

logger: logging.Logger = logging.getLogger(__name__)
PROJECT_DIR: Path = Path(__file__).resolve().parent.parent.parent

def get_cache_status_text() -> str:
    """Формирует строку общей статистики для верхнего табло UI."""
    from layout_generator.assets import _ASSET_CACHE
    products: dict = _ASSET_CACHE.get('products', {})
    if not products:
        return "🔴 Кэш пуст. Запустите первую генерацию для прогрева сервера."
    
    categories: set = set()
    sku_count: int = 0
    for key in products.keys():
        if key.startswith('products_hierarchy.'): continue
        parts = key.split('.')
        if len(parts) >= 3:
            categories.add(parts[1].upper())
            sku_count += 1
            
    return (f"🟢 СЕРВЕР АКТИВЕН. Доступно SKU: {sku_count} | Категорий: {len(categories)}\n"
            f"Загружено оборудования: 4 модели (Холодильники и Стеллажи)")

def get_inventory_json() -> Dict[str, Any]:
    """Формирует дерево товаров для JSON-инспектора."""
    from layout_generator.assets import _ASSET_CACHE
    products: dict = _ASSET_CACHE.get('products', {})
    if not products: return {}
    
    tree: Dict[str, List[str]] = {}
    for key in products.keys():
        if key.startswith('products_hierarchy.'): continue
        parts = key.split('.')
        if len(parts) >= 3:
            cat = parts[1].upper()
            item = ".".join(parts[2:])
            if cat not in tree: tree[cat] = []
            if item not in tree[cat]: tree[cat].append(item)
            
    return {k: sorted(v) for k, v in sorted(tree.items())}

def format_checklist(steps_state: List[Dict[str, Any]]) -> str:
    """Генерирует текст чек-листа прогресса."""
    lines: List[str] = ["🚀 Процесс генерации...\n"]
    now: float = time.time()
    for s in steps_state:
        status: str = s["status"]
        if status == "pending": lines.append(f"[  ] ⚪ {s['name']}")
        elif status == "running":
            el = now - s.get("start_time", now)
            lines.append(f"[⏳] 🟡 {s['name']}... ({el:.1f} сек)")
        elif status == "done":
            lines.append(f"[✅] 🟢 {s['name']} ({s.get('time', 0):.1f} сек)")
        elif status == "error":
            lines.append(f"[❌] 🔴 {s['name']} - {s.get('message')}")
    return "\n".join(lines)

def generate_wrapper(user_prompt: str) -> Generator[
    Tuple[
        Optional[str],  # 1. output_3d (Путь к GLB модели)
        str,            # 2. status_output (Текст чек-листа)
        Optional[str],  # 3. output_svg (Путь к картинке для превью 2D-плана)
        Optional[str],  # 4. draft_json_out (Путь к layout_draft.json)
        Optional[str],  # 5. final_json_out (Путь к layout_final.json)
        Optional[str],  # 6. svg_file_out (Путь к layout_final.svg для скачивания)
        Optional[str],  # 7. history_json_out (Путь к llm_history.jsonl)
        Optional[str]   # 8. draft_svg_file_out (Draft SVG)
    ], 
    None, 
    None
]:
    """Оркестратор пайплайна с передачей промпта в LLM и запуском 3D-движка."""
    tracker = setup_logger()
    local_logger = logging.getLogger(__name__)

    # Чек-лист остается без изменений...
    steps_state: List[Dict[str, Any]] = [
        {"name": "ИИ-анализ и граф", "status": "pending"},
        {"name": "Инициализация", "status": "pending"},
        {"name": "Загрузка ассетов", "status": "pending"},
        {"name": "Расчет планограммы", "status": "pending"},
        {"name": "Генерация топологии", "status": "pending"},
        {"name": "Сборка 3D-сцены", "status": "pending"},
        {"name": "Экспорт GLB (Draco)", "status": "pending"}
    ]

    # Важно: yield теперь должен возвращать 8 элементов (None для всех файлов, пока идет генерация)
    yield None, format_checklist(steps_state), None, None, None, None, None, None
    
    total_start: float = time.time()
    task_queue: queue.Queue = queue.Queue()

    def worker_thread() -> None:
        try:
            main_cfg = OmegaConf.load(PROJECT_DIR / "configs" / "main_config.yaml")
            layout_cfg = OmegaConf.load(PROJECT_DIR / "configs" / "layout_config.yaml")
            asset_cfg = OmegaConf.load(PROJECT_DIR / "configs" / "asset_config.yaml")
            
            llm_graph: Optional[Dict[str, Any]] = None

            if user_prompt.strip():
                task_queue.put({"type": "step_update", "idx": 0, "update": {"status": "running"}})
                llm_start: float = time.time()
                
                asset_cfg_raw = OmegaConf.to_container(asset_cfg, resolve=False)
                
                categories: List[str] = []
                for group in asset_cfg_raw.get("products_hierarchy", {}).values():
                    if isinstance(group, dict):
                        categories.extend(list(group.keys()))
                
                equipment_lines: List[str] = []
                fixtures_cfg = asset_cfg_raw.get("fixtures", asset_cfg_raw.get("equipment", {})) 
                base_assets_dir = asset_cfg_raw.get("assets_dir_path", "assets")
                
                for fix_id, fix_data in fixtures_cfg.items():
                    if not isinstance(fix_data, dict): continue
                    raw_path = fix_data.get("asset_file_path", "")
                    if not raw_path: continue
                    clean_path = raw_path.replace("${assets.assets_dir_path}", base_assets_dir)
                    full_path = PROJECT_DIR / clean_path
                    w, d = None, None
                    try:
                        import trimesh
                        scene = trimesh.load(str(full_path), force='mesh')
                        w = round(float(scene.extents[0]), 2)
                        d = round(float(scene.extents[1]), 2)
                    except Exception as e:
                        local_logger.debug(f"⚠️ Не удалось вычислить габариты для {fix_id}: {e}")
                    
                    if w is not None and d is not None:
                        desc: str = fix_data.get("asset_name", fix_id)
                        equipment_lines.append(f"- {fix_id} ({desc}): {w}m x {d}m")
                
                equipment_catalog: str = "\n".join(equipment_lines)
                
                if not equipment_catalog:
                    equipment_catalog = (
                        "- showcase_glb (Холодильник пристенный): 1.2m x 0.6m\n"
                        "- shelf_metal (Стеллаж высокий): 1.0m x 0.4m\n"
                        "- small_shelf_two_sided (Островок двусторонний): 1.2m x 0.8m"
                    )

                client = InventoryLLMClient(main_cfg)
                result = client.generate_layout(
                    user_request=user_prompt, 
                    available_categories=categories,
                    equipment_catalog=equipment_catalog,
                    size_n=layout_cfg.layout.size_n,
                    size_m=layout_cfg.layout.size_m
                )                
                if result:
                    llm_graph = result
                    try:
                        from layout_generator.visualizer import draw_llm_draft_plan
                        draw_llm_draft_plan(llm_graph, PROJECT_DIR)
                    except Exception as ve:
                        local_logger.error(f"Визуализация не удалась: {ve}")
                
                task_queue.put({"type": "step_update", "idx": 0, "update": {"status": "done", "time": time.time() - llm_start}})
            else:
                task_queue.put({"type": "step_update", "idx": 0, "update": {"status": "done", "time": 0.0}})

            for res_path, update in generate_layout_stream(
                project_dir=PROJECT_DIR, 
                tracker=tracker, 
                size_n=layout_cfg.layout.size_n, 
                size_m=layout_cfg.layout.size_m,
                layout_graph=llm_graph
            ):
                if update and "step_idx" in update:
                    task_queue.put({"type": "step_update", "idx": update["step_idx"], "update": update})
                if res_path:
                    task_queue.put({"type": "finish", "result_path": res_path})
                    return
        except Exception as e:
            local_logger.error(f"Ошибка воркера: {e}", exc_info=True)
            task_queue.put({"type": "error", "message": str(e)})

    thread = threading.Thread(target=worker_thread)
    thread.start()

    res_path: Optional[str] = None
    
    while thread.is_alive() or not task_queue.empty():
        try:
            msg = task_queue.get(timeout=0.1)
            if msg["type"] == "step_update":
                idx, upd = msg["idx"], msg["update"]
                if upd.get("status") == "running" and "start_time" not in steps_state[idx]:
                    steps_state[idx]["start_time"] = time.time()
                steps_state[idx].update(upd)
                yield None, format_checklist(steps_state), None, None, None, None, None, None
            elif msg["type"] == "finish": 
                res_path = msg["result_path"]
            elif msg["type"] == "error": 
                raise gr.Error(msg["message"])
        except queue.Empty:
            yield None, format_checklist(steps_state), None, None, None, None, None, None

    if res_path:
        final_text: str = format_checklist(steps_state)
        final_text += "\n" + "-"*40 + "\n"
        final_text += f"🏁 Бэкенд отработал успешно! Время: {time.time() - total_start:.1f} сек.\n"
        final_text += "⏳ Браузер скачивает артефакты..."
        
        absolute_model_path: str = str(Path(res_path).resolve())
        
        # === Обновленный сбор путей к артефактам (всё из папки logs) ===
        logs_dir = PROJECT_DIR / "logs"
        
        # Ищем JSON
        draft_path: Optional[str] = str(logs_dir / "layout_draft.json") if (logs_dir / "layout_draft.json").exists() else None
        final_path: Optional[str] = str(logs_dir / "layout_final.json") if (logs_dir / "layout_final.json").exists() else None
        history_path: Optional[str] = str(logs_dir / "llm_history.jsonl") if (logs_dir / "llm_history.jsonl").exists() else None
        
        # Ищем SVG тоже в logs
        svg_path: Optional[str] = str(logs_dir / "layout_final.svg") if (logs_dir / "layout_final.svg").exists() else None
        draft_svg_path: Optional[str] = str(logs_dir / "layout_draft.svg") if (logs_dir / "layout_draft.svg").exists() else None
        
        # Возвращаем: 3D_модель, Текст, SVG_превью, Draft_JSON, Final_JSON, SVG_файл, History_JSONL, Draft_SVG
        yield absolute_model_path, final_text, svg_path, draft_path, final_path, svg_path, history_path, draft_svg_path