"""
engine/training_setup.py
Manages the auto-spawned ModelBlock and DataLoaderBlock nodes in the Training tab.

"""

from __future__ import annotations

import dearpygui.dearpygui as dpg

import ml_forge.state as state
from ml_forge.ui.console import log

_MODEL_POS      = (40,  60)
_DATALOADER_POS = (40, 260)
_MODEL_NID      = 9901
_DATALOADER_NID = 9902


def _get_training_tab() -> dict | None:
    for t in state.tabs.values():
        if t.get("role") == "training":
            return t
    return None


def _tid_of(tab: dict) -> int | None:
    for tid, t in state.tabs.items():
        if t is tab:
            return tid
    return None


def _ntag(tid: int, nid: int) -> str:
    return f"node_{tid}_{nid}"


def ensure_pipeline_inputs() -> None:
    tab = _get_training_tab()
    if tab is None:
        return
    tid = _tid_of(tab)
    if tid is None:
        return

    from ml_forge.graph.nodes import raw_spawn_node

    model_tag  = _ntag(tid, _MODEL_NID)
    loader_tag = _ntag(tid, _DATALOADER_NID)

    if not dpg.does_item_exist(model_tag):
        raw_spawn_node(tid, "ModelBlock", nid=_MODEL_NID, pos=_MODEL_POS)
        _lock_node(model_tag)

    if not dpg.does_item_exist(loader_tag):
        raw_spawn_node(tid, "DataLoaderBlock", nid=_DATALOADER_NID, pos=_DATALOADER_POS)
        _lock_node(loader_tag)


def _lock_node(ntag: str) -> None:
    if not dpg.does_item_exist(ntag):
        return
    with dpg.theme() as th:
        with dpg.theme_component(dpg.mvNode):
            dpg.add_theme_color(dpg.mvNodeCol_TitleBar,         (50, 50, 80),   category=dpg.mvThemeCat_Nodes)
            dpg.add_theme_color(dpg.mvNodeCol_TitleBarHovered,  (70, 70, 110),  category=dpg.mvThemeCat_Nodes)
            dpg.add_theme_color(dpg.mvNodeCol_TitleBarSelected, (90, 90, 140),  category=dpg.mvThemeCat_Nodes)
            dpg.add_theme_color(dpg.mvNodeCol_NodeOutline,      (80, 120, 200), category=dpg.mvThemeCat_Nodes)
    dpg.bind_item_theme(ntag, th)


def update_block_labels(epoch: int = 0, total: int = 0,
                         train_loss: float | None = None,
                         val_loss: float | None = None,
                         val_acc: float | None = None) -> None:
    """
    Update ModelBlock and DataLoaderBlock node titles with live training stats.
    Called from drain_result_queue() on each epoch result.

    DearPyGui nodes don't expose a title text item directly — we update the
    node label via dpg.set_item_label(), which changes the title bar text.
    """
    tab = _get_training_tab()
    if tab is None:
        return
    tid = _tid_of(tab)
    if tid is None:
        return

    model_tag  = _ntag(tid, _MODEL_NID)
    loader_tag = _ntag(tid, _DATALOADER_NID)

    # Build label strings
    if epoch > 0 and total > 0:
        progress = f"[{epoch}/{total}]"
    else:
        progress = ""

    if train_loss is not None:
        loss_str = f"loss={train_loss:.3f}"
        if val_loss is not None:
            loss_str += f"  val={val_loss:.3f}"
        if val_acc is not None:
            loss_str += f"  acc={val_acc:.2%}"
        model_label  = f"ModelBlock  {progress}  {loss_str}"
        loader_label = f"DataLoaderBlock  {progress}"
    else:
        model_label  = f"ModelBlock  {progress}"
        loader_label = f"DataLoaderBlock  {progress}"

    if dpg.does_item_exist(model_tag):
        try:
            dpg.set_item_label(model_tag, model_label)
        except Exception:
            pass

    if dpg.does_item_exist(loader_tag):
        try:
            dpg.set_item_label(loader_tag, loader_label)
        except Exception:
            pass


def reset_block_labels() -> None:
    """
    Restore ModelBlock and DataLoaderBlock node titles to their original labels.
    Called when training stops or finishes.
    """
    tab = _get_training_tab()
    if tab is None:
        return
    tid = _tid_of(tab)
    if tid is None:
        return

    model_tag  = _ntag(tid, _MODEL_NID)
    loader_tag = _ntag(tid, _DATALOADER_NID)

    if dpg.does_item_exist(model_tag):
        try:
            dpg.set_item_label(model_tag, "ModelBlock")
        except Exception:
            pass

    if dpg.does_item_exist(loader_tag):
        try:
            dpg.set_item_label(loader_tag, "DataLoaderBlock")
        except Exception:
            pass