gui.py

#!/usr/bin/env python3
"""
TL++ GUI Application (PyQt6)

PyQt6-based GUI for the TL++ distributed federated learning system.
Supports Orchestrator, Helper, and Node roles with real-time log display.

Usage:
    pip3 install PyQt6
    python3 gui.py
"""

import sys
import re
import time
import json
import threading
import queue
import logging
import socket
import traceback
from pathlib import Path
import platform
import webbrowser
from datetime import datetime, timezone, timedelta


from docker_manager import DockerManager, should_use_docker

from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QGroupBox, QLabel, QLineEdit, QSpinBox, QDoubleSpinBox,
    QCheckBox, QPushButton, QTextEdit, QMessageBox, QFrame,
    QScrollArea, QProgressBar, QGridLayout, QDialog, QFileDialog
)
from PyQt6.QtCore import Qt, QTimer, pyqtSignal, QThread
from PyQt6.QtGui import QTextCursor, QColor

# Docker 관리 모듈 import
try:
    from docker_manager import DockerManager, DockerStatusChecker, should_use_docker
    DOCKER_AVAILABLE = True
except ImportError:
    DOCKER_AVAILABLE = False
    def should_use_docker():
        return platform.system() in ('Darwin', 'Windows')


# ==============================================================================
# CONSTANTS
# ==============================================================================

KST = timezone(timedelta(hours=9))

DEFAULT_ORCH_PORT = 8080
DEFAULT_HELPER_NODE_PORT = 8081
DEFAULT_HELPER_ORCH_PORT = 8082

CONFIG_DIR = Path.home() / '.tlpp'
LAST_CONFIG_PATH = CONFIG_DIR / 'last_config.json'

# Modern Color Palette - Dark Theme
COLORS_DARK = {
    'bg_dark': '#121212',
    'bg_card': '#1e1e1e',
    'bg_input': '#2d2d2d',
    'accent_primary': '#ff4757',
    'accent_secondary': '#3d3d3d',
    'accent_success': '#2ed573',
    'accent_warning': '#ffa502',
    'accent_info': '#3742fa',
    'text_primary': '#ffffff',
    'text_secondary': '#b0b0b0',
    'text_muted': '#707070',
    'border': '#404040',
}

# Modern Color Palette - Light Theme
COLORS_LIGHT = {
    'bg_dark': '#f5f5f5',
    'bg_card': '#ffffff',
    'bg_input': '#e8e8e8',
    'accent_primary': '#e63946',
    'accent_secondary': '#d0d0d0',
    'accent_success': '#2d9a46',
    'accent_warning': '#e67e00',
    'accent_info': '#2563eb',
    'text_primary': '#1a1a1a',
    'text_secondary': '#4a4a4a',
    'text_muted': '#808080',
    'border': '#c0c0c0',
}

# Current theme (default: dark)
COLORS = COLORS_DARK.copy()

def get_stylesheet(colors):
    """Generate stylesheet with given color scheme."""
    return f"""
QMainWindow {{
    background-color: {colors['bg_dark']};
}}

QWidget {{
    background-color: transparent;
    color: {colors['text_primary']};
    font-family: -apple-system, BlinkMacSystemFont, 'Noto Sans KR', 'Malgun Gothic', sans-serif;
}}

QGroupBox {{
    background-color: {colors['bg_card']};
    border: 1px solid {colors['border']};
    border-radius: 12px;
    margin-top: 12px;
    padding: 15px;
    padding-top: 25px;
    font-weight: bold;
    font-size: 13px;
}}

QGroupBox::title {{
    subcontrol-origin: margin;
    subcontrol-position: top left;
    left: 15px;
    padding: 0 8px;
    color: {colors['accent_info']};
}}

QLabel {{
    color: {colors['text_secondary']};
    font-size: 12px;
}}

QLineEdit {{
    background-color: {colors['bg_input']};
    border: 2px solid {colors['border']};
    border-radius: 8px;
    padding: 12px 14px;
    color: {colors['text_primary']};
    font-size: 14px;
    min-height: 20px;
    selection-background-color: {colors['accent_primary']};
}}

QLineEdit:focus {{
    border: 2px solid {colors['accent_primary']};
}}

QLineEdit:disabled {{
    background-color: {colors['bg_dark']};
    color: {colors['text_muted']};
}}

QSpinBox, QDoubleSpinBox {{
    background-color: {colors['bg_input']};
    border: 2px solid {colors['border']};
    border-radius: 8px;
    padding: 10px 12px;
    color: {colors['text_primary']};
    font-size: 14px;
    min-height: 20px;
}}

QSpinBox:focus, QDoubleSpinBox:focus {{
    border: 2px solid {colors['accent_primary']};
}}

QSpinBox:disabled, QDoubleSpinBox:disabled {{
    background-color: {colors['bg_dark']};
    color: {colors['text_muted']};
}}

QSpinBox::up-button, QDoubleSpinBox::up-button {{
    subcontrol-origin: padding;
    subcontrol-position: top right;
    width: 18px;
    border-left: 1px solid {colors['border']};
    border-top-right-radius: 6px;
    background-color: {colors['bg_card']};
}}

QSpinBox::up-button:hover, QDoubleSpinBox::up-button:hover {{
    background-color: {colors['accent_secondary']};
}}

QSpinBox::down-button, QDoubleSpinBox::down-button {{
    subcontrol-origin: padding;
    subcontrol-position: bottom right;
    width: 18px;
    border-left: 1px solid {colors['border']};
    border-bottom-right-radius: 6px;
    background-color: {colors['bg_card']};
}}

QSpinBox::down-button:hover, QDoubleSpinBox::down-button:hover {{
    background-color: {colors['accent_secondary']};
}}

QCheckBox {{
    color: {colors['text_secondary']};
    font-size: 12px;
    spacing: 8px;
}}

QCheckBox::indicator {{
    width: 20px;
    height: 20px;
    border-radius: 6px;
    border: 2px solid {colors['border']};
    background-color: {colors['bg_input']};
}}

QCheckBox::indicator:checked {{
    background-color: {colors['accent_primary']};
    border: 2px solid {colors['accent_primary']};
}}

QCheckBox::indicator:hover {{
    border: 2px solid {colors['accent_primary']};
}}

QRadioButton {{
    color: {colors['text_primary']};
    font-size: 13px;
    spacing: 8px;
}}

QRadioButton::indicator {{
    width: 20px;
    height: 20px;
    border-radius: 10px;
    border: 2px solid {colors['border']};
    background-color: {colors['bg_input']};
}}

QRadioButton::indicator:checked {{
    background-color: {colors['accent_primary']};
    border: 2px solid {colors['accent_primary']};
}}

QRadioButton::indicator:hover {{
    border: 2px solid {colors['accent_primary']};
}}

QPushButton {{
    background-color: {colors['accent_primary']};
    color: white;
    border: none;
    border-radius: 10px;
    padding: 12px 24px;
    font-size: 14px;
    font-weight: bold;
}}

QPushButton:hover {{
    background-color: #ff6b6b;
}}

QPushButton:pressed {{
    background-color: #c73e54;
}}

QPushButton:disabled {{
    background-color: {colors['text_muted']};
    color: {colors['text_secondary']};
}}

QPushButton#stopBtn {{
    background-color: {colors['bg_input']};
    border: 2px solid {colors['accent_primary']};
    color: {colors['accent_primary']};
}}

QPushButton#stopBtn:hover {{
    background-color: {colors['accent_primary']};
    color: white;
}}

QPushButton#stopBtn:disabled {{
    background-color: {colors['bg_input']};
    border: 2px solid {colors['text_muted']};
    color: {colors['text_muted']};
}}

QPushButton#clearBtn {{
    background-color: transparent;
    border: 1px solid {colors['border']};
    color: {colors['text_secondary']};
    padding: 6px 12px;
    font-size: 11px;
}}

QPushButton#clearBtn:hover {{
    border: 1px solid {colors['accent_info']};
    color: {colors['accent_info']};
}}

QPushButton#themeBtn {{
    background-color: {colors['bg_input']};
    border: 1px solid {colors['border']};
    color: {colors['text_secondary']};
    padding: 8px 12px;
    font-size: 16px;
    border-radius: 8px;
    min-width: 40px;
}}

QPushButton#themeBtn:hover {{
    border: 1px solid {colors['accent_info']};
    color: {colors['accent_info']};
}}

QTextEdit {{
    background-color: {colors['bg_dark']};
    border: 2px solid {colors['border']};
    border-radius: 10px;
    padding: 10px;
    color: {colors['text_primary']};
    font-family: 'JetBrains Mono', 'Fira Code', 'Consolas', monospace;
    font-size: 11px;
    selection-background-color: {colors['accent_primary']};
}}

QFrame#statusFrame {{
    background-color: {colors['bg_card']};
    border: 1px solid {colors['border']};
    border-radius: 10px;
    padding: 10px;
}}

QScrollBar:vertical {{
    background-color: {colors['bg_dark']};
    width: 12px;
    border-radius: 6px;
}}

QScrollBar::handle:vertical {{
    background-color: {colors['border']};
    border-radius: 6px;
    min-height: 30px;
}}

QScrollBar::handle:vertical:hover {{
    background-color: {colors['accent_primary']};
}}

QScrollBar::add-line:vertical, QScrollBar::sub-line:vertical {{
    height: 0px;
}}

QProgressBar {{
    background-color: {colors['bg_input']};
    border: 1px solid {colors['border']};
    border-radius: 8px;
    text-align: center;
    color: {colors['text_primary']};
    font-size: 12px;
    font-weight: bold;
    min-height: 22px;
    max-height: 22px;
}}

QProgressBar::chunk {{
    background-color: {colors['accent_primary']};
    border-radius: 7px;
}}

QProgressBar#batchBar {{
    min-height: 8px;
    max-height: 8px;
    border-radius: 4px;
    font-size: 0px;
}}

QProgressBar#batchBar::chunk {{
    background-color: {colors['accent_warning']};
    border-radius: 3px;
}}
"""

# Default stylesheet (dark theme)
STYLESHEET = get_stylesheet(COLORS)


# ==============================================================================
# LOGGING
# ==============================================================================

class KSTFormatter(logging.Formatter):
    """Custom log formatter using Korean Standard Time."""

    def formatTime(self, record, datefmt=None):
        dt = datetime.fromtimestamp(record.created, tz=KST)
        if datefmt:
            return dt.strftime(datefmt)
        return dt.isoformat()


class QueueHandler(logging.Handler):
    """Routes logging to queue for thread-safe GUI updates."""

    def __init__(self, log_queue):
        super().__init__()
        self.queue = log_queue
        self.setFormatter(KSTFormatter(
            fmt='%(asctime)s | %(levelname)-8s | %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S KST'
        ))

    def emit(self, record):
        msg = self.format(record)
        self.queue.put(('log', record.levelname, msg))


class StderrCapture:
    """Captures stderr (tqdm output) and routes to message queue."""

    def __init__(self, message_queue):
        self.queue = message_queue
        self.original_stderr = sys.stderr
        self.buffer = ""

    def write(self, text):
        # tqdm uses \r for progress updates, filter and clean up
        if text:
            # Strip carriage returns and clean up
            clean_text = text.replace('\r', '').strip()
            if clean_text and not clean_text.isspace():
                self.queue.put(('log', 'INFO', clean_text))
        return len(text)

    def flush(self):
        pass

    def restore(self):
        sys.stderr = self.original_stderr


# ==============================================================================
# UTILITY FUNCTIONS
# ==============================================================================

def get_local_ip_addresses():
    """Get local IP addresses for display."""
    ips = []
    try:
        hostname = socket.gethostname()
        ips.append(socket.gethostbyname(hostname))
    except:
        pass

    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        if ip not in ips:
            ips.append(ip)
        s.close()
    except:
        pass

    if not ips:
        ips = ['127.0.0.1']

    return ips


# ==============================================================================
# DOCKER SETUP DIALOG
# ==============================================================================

class DockerSetupDialog(QDialog):
    """Docker 설치 안내 다이얼로그"""

    def __init__(self, parent=None, download_url: str = ""):
        super().__init__(parent)
        self.download_url = download_url
        self.setWindowTitle("환경 설정 필요")
        self.setFixedSize(450, 280)
        self.setModal(True)

        self._setup_ui()

    def _setup_ui(self):
        layout = QVBoxLayout(self)
        layout.setSpacing(20)
        layout.setContentsMargins(30, 30, 30, 30)

        # 아이콘 + 제목
        title_layout = QHBoxLayout()
        icon_label = QLabel("🐳")
        icon_label.setStyleSheet("font-size: 48px;")
        title_layout.addWidget(icon_label)

        title_text = QLabel("Docker 설치가 필요합니다")
        title_text.setStyleSheet(f"font-size: 18px; font-weight: bold; color: {COLORS['text_primary']};")
        title_layout.addWidget(title_text)
        title_layout.addStretch()
        layout.addLayout(title_layout)

        # 설명
        desc_label = QLabel(
            "TL++를 실행하려면 Docker Desktop이 필요합니다.\n"
            "Docker는 복잡한 환경 설정 없이 TL++를 실행할 수 있게 해주는\n"
            "도구입니다. 아래 버튼을 클릭하여 설치해주세요."
        )
        desc_label.setStyleSheet(f"font-size: 13px; color: {COLORS['text_secondary']}; line-height: 1.5;")
        desc_label.setWordWrap(True)
        layout.addWidget(desc_label)

        layout.addStretch()

        # 버튼들
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(15)

        download_btn = QPushButton("Docker 다운로드")
        download_btn.setStyleSheet(f"""
            QPushButton {{
                background-color: {COLORS['accent_info']};
                color: white;
                border: none;
                border-radius: 8px;
                padding: 12px 24px;
                font-size: 14px;
                font-weight: bold;
            }}
            QPushButton:hover {{
                background-color: #4a5aff;
            }}
        """)
        download_btn.clicked.connect(self._open_download)
        btn_layout.addWidget(download_btn)

        retry_btn = QPushButton("다시 확인")
        retry_btn.setStyleSheet(f"""
            QPushButton {{
                background-color: {COLORS['bg_input']};
                color: {COLORS['text_primary']};
                border: 2px solid {COLORS['border']};
                border-radius: 8px;
                padding: 12px 24px;
                font-size: 14px;
            }}
            QPushButton:hover {{
                border-color: {COLORS['accent_primary']};
            }}
        """)
        retry_btn.clicked.connect(self.accept)
        btn_layout.addWidget(retry_btn)

        cancel_btn = QPushButton("취소")
        cancel_btn.setStyleSheet(f"""
            QPushButton {{
                background-color: transparent;
                color: {COLORS['text_muted']};
                border: none;
                padding: 12px 16px;
                font-size: 13px;
            }}
            QPushButton:hover {{
                color: {COLORS['text_primary']};
            }}
        """)
        cancel_btn.clicked.connect(self.reject)
        btn_layout.addWidget(cancel_btn)

        layout.addLayout(btn_layout)

    def _open_download(self):
        """Docker 다운로드 페이지 열기"""
        if self.download_url:
            webbrowser.open(self.download_url)


class DockerNotRunningDialog(QDialog):
    """Docker가 설치되었지만 실행 중이 아닐 때 표시하는 다이얼로그"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Docker 시작 필요")
        self.setFixedSize(400, 220)
        self.setModal(True)

        self._setup_ui()

    def _setup_ui(self):
        layout = QVBoxLayout(self)
        layout.setSpacing(20)
        layout.setContentsMargins(30, 30, 30, 30)

        # 아이콘 + 제목
        title_layout = QHBoxLayout()
        icon_label = QLabel("⏳")
        icon_label.setStyleSheet("font-size: 40px;")
        title_layout.addWidget(icon_label)

        title_text = QLabel("Docker를 시작해주세요")
        title_text.setStyleSheet(f"font-size: 18px; font-weight: bold; color: {COLORS['text_primary']};")
        title_layout.addWidget(title_text)
        title_layout.addStretch()
        layout.addLayout(title_layout)

        # 설명
        desc_label = QLabel(
            "Docker Desktop이 설치되어 있지만 실행 중이 아닙니다.\n"
            "Docker Desktop을 시작한 후 '다시 확인' 버튼을 클릭해주세요."
        )
        desc_label.setStyleSheet(f"font-size: 13px; color: {COLORS['text_secondary']};")
        desc_label.setWordWrap(True)
        layout.addWidget(desc_label)

        layout.addStretch()

        # 버튼
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(15)

        retry_btn = QPushButton("다시 확인")
        retry_btn.clicked.connect(self.accept)
        btn_layout.addWidget(retry_btn)

        cancel_btn = QPushButton("취소")
        cancel_btn.setObjectName("clearBtn")
        cancel_btn.clicked.connect(self.reject)
        btn_layout.addWidget(cancel_btn)

        layout.addLayout(btn_layout)


# ==============================================================================
# IMAGE BUILD DIALOG
# ==============================================================================

class ImageBuildThread(QThread):
    """Docker 이미지 빌드를 위한 백그라운드 스레드"""

    progress = pyqtSignal(str, int)  # message, percent
    finished = pyqtSignal(bool)  # success

    def __init__(self, docker_manager):
        super().__init__()
        self.docker_manager = docker_manager

    def run(self):
        success = self.docker_manager.build_image(
            progress_callback=lambda msg, pct: self.progress.emit(msg, pct)
        )
        self.finished.emit(success)


class ImageBuildDialog(QDialog):
    """Docker 이미지 빌드 프로그레스 다이얼로그"""

    def __init__(self, docker_manager, parent=None):
        super().__init__(parent)
        self.docker_manager = docker_manager
        self.build_success = False

        self.setWindowTitle("환경 준비")
        self.setFixedSize(450, 200)
        self.setModal(True)

        # 닫기 버튼 비활성화 (빌드 중 강제 종료 방지)
        self.setWindowFlags(
            self.windowFlags() & ~Qt.WindowType.WindowCloseButtonHint
        )

        self._setup_ui()

    def _setup_ui(self):
        layout = QVBoxLayout(self)
        layout.setSpacing(20)
        layout.setContentsMargins(30, 30, 30, 30)

        # 제목
        self.title_label = QLabel("환경을 준비하고 있습니다...")
        self.title_label.setStyleSheet(f"font-size: 16px; font-weight: bold; color: {COLORS['text_primary']};")
        self.title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.title_label)

        # 상태 메시지
        self.status_label = QLabel("Docker 이미지를 빌드하고 있습니다...")
        self.status_label.setStyleSheet(f"font-size: 12px; color: {COLORS['text_secondary']};")
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.status_label)

        # 프로그레스 바
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setTextVisible(True)
        layout.addWidget(self.progress_bar)

        # 안내 메시지
        hint_label = QLabel("처음 실행 시에만 시간이 걸립니다. 잠시만 기다려주세요.")
        hint_label.setStyleSheet(f"font-size: 11px; color: {COLORS['text_muted']};")
        hint_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(hint_label)

        layout.addStretch()

    def start_build(self):
        """빌드 시작"""
        self.build_thread = ImageBuildThread(self.docker_manager)
        self.build_thread.progress.connect(self._on_progress)
        self.build_thread.finished.connect(self._on_finished)
        self.build_thread.start()

    def _on_progress(self, message: str, percent: int):
        """빌드 진행 상황 업데이트"""
        self.status_label.setText(message)
        self.progress_bar.setValue(percent)

    def _on_finished(self, success: bool):
        """빌드 완료"""
        self.build_success = success
        if success:
            self.title_label.setText("준비 완료!")
            self.status_label.setText("환경이 성공적으로 준비되었습니다.")
            self.progress_bar.setValue(100)
            # 잠시 후 자동 닫기
            QTimer.singleShot(1000, self.accept)
        else:
            self.title_label.setText("준비 실패")
            self.status_label.setText("환경 준비에 실패했습니다. 다시 시도해주세요.")
            self.progress_bar.setValue(0)
            # 닫기 버튼 활성화
            self.setWindowFlags(
                self.windowFlags() | Qt.WindowType.WindowCloseButtonHint
            )
            self.show()


# ==============================================================================
# TRAINING THREAD
# ==============================================================================

class TrainingThread(threading.Thread):
    """Background thread for running training processes."""

    def __init__(self, role, config, message_queue, use_docker=False, docker_manager=None):
        super().__init__(daemon=True)
        self.role = role
        self.config = config
        self.message_queue = message_queue
        self.stop_event = threading.Event()
        self.instance = None
        self.use_docker = use_docker
        self.docker_manager = None
        self.docker_container_name = None

    def run(self):
        self._setup_logging()

        try:
            if self.use_docker and self.docker_manager:
                self._run_docker()
            elif self.role == 'orchestrator':
                self._run_orchestrator()
            elif self.role == 'helper':
                self._run_helper()
            else:
                self._run_node()
        except Exception as e:
            logging.error(f"Failed: {e}")
            logging.error(traceback.format_exc())
        finally:
            self._cleanup()
            self.message_queue.put(('status', 'stopped', None))

    def _setup_logging(self):
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        logger.handlers.clear()

        handler = QueueHandler(self.message_queue)
        handler.setLevel(logging.INFO)
        logger.addHandler(handler)

        # Capture stderr for tqdm progress bars
        self.stderr_capture = StderrCapture(self.message_queue)
        sys.stderr = self.stderr_capture

    def _run_orchestrator(self):
        from TLpp.orchestrator import SecureOrchestrator

        self.message_queue.put(('status', 'setup', 'Orchestrator starting...'))
        logging.info("=" * 60)
        logging.info("TL++ Orchestrator")
        logging.info("=" * 60)

        self.instance = SecureOrchestrator(self.config)

        logging.info(f"Secure mode: {self.config.get('secure')}")
        logging.info(f"Waiting for {self.config.get('n_nodes')} node(s)...")

        self.instance.setup()
        self.message_queue.put(('status', 'running', 'Training in progress'))

        self.instance.train()

        logging.info("Training completed")
        self.message_queue.put(('status', 'completed', 'Training completed'))

    def _run_node(self):
        self.message_queue.put(('status', 'setup', 'Connecting to orchestrator...'))
        logging.info("=" * 60)

        if should_use_docker():
            self._run_node_docker()
        else:
            self._run_node_native()

    def _run_node_docker(self):
        """Mac/Windows: Docker 컨테이너에서 Node 실행."""
        import platform
        logging.info(f"TL++ Training Node (Docker Mode - {platform.system()})")
        logging.info("=" * 60)

        self.docker_manager = DockerManager()

        # Docker 설치/실행 확인
        if not self.docker_manager.is_docker_installed():
            url = self.docker_manager.get_docker_download_url()
            logging.error(f"Docker가 설치되어 있지 않습니다. Docker Desktop을 설치해주세요: {url}")
            self.message_queue.put(('status', 'error', 'Docker not installed'))
            return

        if not self.docker_manager.is_docker_running():
            logging.error("Docker가 실행 중이 아닙니다. Docker Desktop을 시작해주세요.")
            self.message_queue.put(('status', 'error', 'Docker not running'))
            return

        # 이미지 빌드 (없으면)
        self.message_queue.put(('status', 'setup', 'Building Docker image...'))

        def build_progress(msg, percent):
            logging.info(f"[Docker Build] {msg} ({percent}%)")

        if not self.docker_manager.build_image(progress_callback=build_progress):
            logging.error("Docker 이미지 빌드 실패")
            self.message_queue.put(('status', 'error', 'Docker build failed'))
            return

        # 컨테이너 실행
        self.message_queue.put(('status', 'running', 'Training in Docker container'))

        def log_callback(level, line):
            self.message_queue.put(('log', level, line))
            # Node ID 파싱
            if 'Assigned Node ID:' in line or 'Assigned node ID:' in line:
                try:
                    node_id = int(line.split(':')[-1].strip().split()[0])
                    self.message_queue.put(('node_id', node_id, None))
                except (ValueError, IndexError):
                    pass

        container_name = self.docker_manager.run_node(
            config=self.config,
            log_callback=log_callback
        )

        if container_name is None:
            logging.error("Docker 컨테이너 실행 실패")
            self.message_queue.put(('status', 'error', 'Container start failed'))
            return

        self.docker_container_name = container_name

        # 컨테이너 종료 대기
        process = self.docker_manager.running_containers.get(container_name)
        if process:
            process.wait()
            if process.returncode == 0:
                logging.info("Training completed")
                self.message_queue.put(('status', 'completed', 'Training completed'))
            else:
                logging.error(f"Docker container exited with code {process.returncode}")
                self.message_queue.put(('status', 'error', 'Training failed'))

    def _run_node_native(self):
        """Linux: 네이티브 모드로 Node 실행."""
        isolated = self.config.get('isolated', False)

        if isolated:
            from TLpp.node import IsolatedNode
            logging.info("TL++ Training Node (Isolated Mode)")
            logging.info("=" * 60)

            self.instance = IsolatedNode(self.config)
            self.message_queue.put(('status', 'running', 'Training in progress'))
            self.instance.setup_and_run()
        else:
            from TLpp.node import SecureNode
            logging.info("TL++ Training Node (Standard Mode)")
            logging.info("=" * 60)

            self.instance = SecureNode(self.config)
            self.instance.setup()

            node_id = getattr(self.instance, 'node_id', None)
            if node_id is not None:
                self.message_queue.put(('node_id', node_id, None))
                logging.info(f"Assigned Node ID: {node_id}")

            self.message_queue.put(('status', 'running', 'Training in progress'))
            self.instance.run()

        logging.info("Training completed")
        self.message_queue.put(('status', 'completed', 'Training completed'))

    def _run_helper(self):
        from TLpp.helper import HelperNode

        self.message_queue.put(('status', 'setup', 'Helper starting...'))
        logging.info("=" * 60)
        logging.info("TL++ Helper Node (Secure Mode)")
        logging.info("=" * 60)

        self.instance = HelperNode(self.config)

        self.instance.start_node_server()
        self.instance.connect_to_orchestrator()

        n_nodes = self.config.get('n_nodes', 1)
        logging.info(f"Waiting for {n_nodes} node(s)...")
        self.message_queue.put(('status', 'waiting', f'Waiting for {n_nodes} node(s)...'))

        self.instance.wait_for_nodes(n_nodes)

        self.message_queue.put(('status', 'running', 'Processing secure MPC'))

        self.instance.run()

        logging.info("Helper completed")
        self.message_queue.put(('status', 'completed', 'Helper completed'))

    def _run_docker(self):
        """Docker 컨테이너로 Node 실행 (Node 전용)"""
        self.message_queue.put(('status', 'setup', 'Node starting (Docker)...'))
        logging.info("=" * 60)
        logging.info("TL++ Node (Docker Mode)")
        logging.info("=" * 60)

        def log_callback(level, message):
            """Docker 로그를 message_queue로 전달"""
            self.message_queue.put(('log', level, message))

        # Node 컨테이너 실행
        self.container_name = self.docker_manager.run_node(
            config=self.config,
            log_callback=log_callback
        )

        if not self.container_name:
            logging.error("Failed to start Docker container")
            self.message_queue.put(('status', 'error', 'Failed to start container'))
            return

        self.message_queue.put(('status', 'running', 'node running'))
        logging.info(f"Container started: {self.container_name}")

        # 컨테이너가 종료될 때까지 대기
        while not self.stop_event.is_set():
            status = self.docker_manager.get_container_status(self.container_name)
            if status is None or status not in ('running', 'created'):
                break
            time.sleep(1)

        logging.info(f"node completed")
        self.message_queue.put(('status', 'completed', 'node completed'))

    def _cleanup(self):
        # Restore stderr
        if hasattr(self, 'stderr_capture'):
            self.stderr_capture.restore()

        # Docker 컨테이너 정리
        if self.docker_manager:
            try:
                self.docker_manager.cleanup()
            except:
                pass

        if self.instance:
            try:
                self.instance.cleanup()
            except:
                pass

    def request_stop(self):
        self.stop_event.set()

        # Docker 컨테이너 중지
        if self.docker_container_name and self.docker_manager:
            try:
                logging.info(f"Stopping Docker container: {self.docker_container_name}")
                self.docker_manager.stop_container(self.docker_container_name)
            except:
                pass

        if self.instance:
            try:
                self.instance.cleanup()
            except:
                pass


# ==============================================================================
# ROLE CARD WIDGET
# ==============================================================================

class RoleCard(QFrame):
    """Clickable role selection card."""

    def __init__(self, role_id, title, description, icon, callback):
        super().__init__()
        self.role_id = role_id
        self.selected = False
        self.callback = callback

        self.setFixedSize(220, 160)
        self.setCursor(Qt.CursorShape.PointingHandCursor)
        self._update_style()

        layout = QVBoxLayout(self)
        layout.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.setSpacing(12)
        layout.setContentsMargins(20, 20, 20, 20)

        # Icon
        self.icon_label = QLabel(icon)
        self.icon_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.icon_label)

        # Title
        self.title_label = QLabel(title)
        self.title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.title_label)

        # Description
        self.desc_label = QLabel(description)
        self.desc_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.desc_label.setWordWrap(True)
        layout.addWidget(self.desc_label)

        self._update_labels()

    def _update_style(self):
        if self.selected:
            self.setStyleSheet(f"""
                RoleCard {{
                    background-color: {COLORS['bg_input']};
                    border: 2px solid {COLORS['accent_primary']};
                    border-radius: 15px;
                }}
            """)
        else:
            self.setStyleSheet(f"""
                RoleCard {{
                    background-color: {COLORS['bg_card']};
                    border: 2px solid {COLORS['border']};
                    border-radius: 15px;
                }}
            """)

    def _update_labels(self):
        self.icon_label.setStyleSheet(f"font-size: 38px; color: {COLORS['text_primary']}; background: transparent; border: none;")
        self.title_label.setStyleSheet(f"font-size: 16px; font-weight: bold; color: {COLORS['text_primary']}; background: transparent; border: none;")
        self.desc_label.setStyleSheet(f"font-size: 11px; color: {COLORS['text_muted']}; background: transparent; border: none;")

    def set_selected(self, selected):
        self.selected = selected
        self._update_style()

    def mousePressEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.callback(self.role_id)
        super().mousePressEvent(event)


# ==============================================================================
# METRIC CARD WIDGET
# ==============================================================================

class MetricCard(QFrame):
    """Small card displaying a metric title and value."""

    def __init__(self, title, initial_value="--", accent_color=None):
        super().__init__()
        self._title = title
        self._accent = accent_color
        self.setMinimumSize(130, 68)

        layout = QVBoxLayout(self)
        layout.setContentsMargins(12, 8, 12, 8)
        layout.setSpacing(2)

        self.title_label = QLabel(title)
        layout.addWidget(self.title_label)

        self.value_label = QLabel(initial_value)
        layout.addWidget(self.value_label)

        self.update_style()

    def set_value(self, value):
        self.value_label.setText(str(value))

    def update_style(self):
        color = self._accent or COLORS['text_primary']
        self.setStyleSheet(f"""
            MetricCard {{
                background-color: {COLORS['bg_input']};
                border: 1px solid {COLORS['border']};
                border-radius: 10px;
            }}
        """)
        self.title_label.setStyleSheet(
            f"font-size: 10px; color: {COLORS['text_muted']}; "
            f"border: none; background: transparent;"
        )
        self.value_label.setStyleSheet(
            f"font-size: 18px; font-weight: bold; color: {color}; "
            f"border: none; background: transparent;"
        )


# ==============================================================================
# MAIN WINDOW
# ==============================================================================

class TLPPMainWindow(QMainWindow):
    """Main application window."""

    def __init__(self):
        super().__init__()

        self.setWindowTitle("AI-Co-Op Distributed Learning Algorithm")
        self.setMinimumSize(900, 800)
        self.resize(950, 850)

        # State
        self.current_role = 'orchestrator'
        self.is_running = False
        self.training_thread = None
        self.role_cards = {}
        self.is_dark_mode = True  # Theme state

        # Training metrics state
        self._training_start_time = None
        self._metric_cards = {}

        # Communication queue
        self.message_queue = queue.Queue()

        # Docker 관련 상태
        self.use_docker = should_use_docker()
        self.docker_manager = None
        self.docker_ready = False

        if DOCKER_AVAILABLE and self.use_docker:
            self.docker_manager = DockerManager()

        # Build UI
        self._create_widgets()

        # Restore last-used config (if any)
        self._restore_last_config()

        # Timer for polling queue
        self.timer = QTimer()
        self.timer.timeout.connect(self._poll_queue)
        self.timer.start(100)

        # Timer for elapsed time display
        self.elapsed_timer = QTimer()
        self.elapsed_timer.timeout.connect(self._update_elapsed_time)
        self.elapsed_timer.setInterval(1000)

        # 시작 시 Docker 상태 확인 (필요한 경우)
        if self.use_docker:
            QTimer.singleShot(500, self._check_docker_on_startup)

    def _create_widgets(self):
        # Create scroll area
        self.scroll_area = QScrollArea()
        self.scroll_area.setWidgetResizable(True)
        self.scroll_area.setFrameShape(QFrame.Shape.NoFrame)
        self.scroll_area.setStyleSheet(f"""
            QScrollArea {{
                background-color: {COLORS['bg_dark']};
                border: none;
            }}
        """)
        self.setCentralWidget(self.scroll_area)

        # Content widget inside scroll area
        self.content_widget = QWidget()
        self.content_widget.setStyleSheet(f"background-color: {COLORS['bg_dark']};")
        self.scroll_area.setWidget(self.content_widget)

        main_layout = QVBoxLayout(self.content_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(30, 30, 30, 30)

        # Header
        self._create_header(main_layout)

        # Role Selection
        self._create_role_selection(main_layout)

        # Configuration
        self._create_configuration(main_layout)

        # Status
        self._create_status(main_layout)

        # Control Buttons
        self._create_controls(main_layout)

        # Training Progress Dashboard (Orchestrator only)
        self._create_training_progress(main_layout)

        # Log Display
        self._create_log_display(main_layout)

        # Add stretch at the end
        main_layout.addStretch()

    def _create_header(self, parent_layout):
        header_frame = QFrame()
        header_layout = QVBoxLayout(header_frame)
        header_layout.setContentsMargins(0, 0, 0, 10)

        # Top row with theme button
        top_row = QHBoxLayout()
        top_row.addStretch()

        # Theme toggle button (top right)
        self.theme_btn = QPushButton("☀️")
        self.theme_btn.setObjectName("themeBtn")
        self.theme_btn.setToolTip("Toggle Light/Dark Mode")
        self.theme_btn.setFixedSize(44, 44)
        self.theme_btn.clicked.connect(self._toggle_theme)
        top_row.addWidget(self.theme_btn)

        header_layout.addLayout(top_row)

        # Title with gradient effect simulation
        self.title_header = QLabel("TL++")
        self.title_header.setStyleSheet(f"""
            font-size: 36px;
            font-weight: bold;
            color: {COLORS['accent_primary']};
        """)
        self.title_header.setAlignment(Qt.AlignmentFlag.AlignCenter)
        header_layout.addWidget(self.title_header)

        self.subtitle_header = QLabel("Distributed Learning System")
        self.subtitle_header.setStyleSheet(f"""
            font-size: 14px;
            color: {COLORS['text_muted']};
            margin-top: -1px;
        """)
        self.subtitle_header.setAlignment(Qt.AlignmentFlag.AlignCenter)
        header_layout.addWidget(self.subtitle_header)

        parent_layout.addWidget(header_frame)

    def _create_role_selection(self, parent_layout):
        group = QGroupBox("Select Your Role")
        layout = QHBoxLayout()
        layout.setSpacing(20)
        layout.setAlignment(Qt.AlignmentFlag.AlignCenter)

        roles = [
            ('orchestrator', 'Orchestrator', 'Central coordinator\nManages training', '🎯'),
            ('helper', 'Helper', 'MPC intermediary\nSecure mode only', '🔐'),
            ('node', 'Node', 'Training worker\nEdge device', '💻'),
        ]

        for role_id, title, desc, icon in roles:
            card = RoleCard(role_id, title, desc, icon, self.select_role)
            if role_id == 'orchestrator':
                card.set_selected(True)
            self.role_cards[role_id] = card
            layout.addWidget(card)

        group.setLayout(layout)
        parent_layout.addWidget(group)

    def select_role(self, role_id):
        self.current_role = role_id
        for rid, card in self.role_cards.items():
            card.set_selected(rid == role_id)
        self._on_role_change(role_id)

    def _create_configuration(self, parent_layout):
        self.config_group = QGroupBox("Configuration")
        self.config_layout = QVBoxLayout()

        # Create config panels for each role
        self._create_orchestrator_config()
        self._create_helper_config()
        self._create_node_config()

        # Show orchestrator by default
        self.helper_config_widget.hide()
        self.node_config_widget.hide()

        # Save / Load preset buttons
        btn_row = QHBoxLayout()
        btn_row.addStretch()
        save_preset_btn = QPushButton("Save Preset")
        save_preset_btn.setObjectName("clearBtn")
        save_preset_btn.clicked.connect(self._on_save_preset)
        btn_row.addWidget(save_preset_btn)
        load_preset_btn = QPushButton("Load Preset")
        load_preset_btn.setObjectName("clearBtn")
        load_preset_btn.clicked.connect(self._on_load_preset)
        btn_row.addWidget(load_preset_btn)
        self.config_layout.addLayout(btn_row)

        self.config_group.setLayout(self.config_layout)
        parent_layout.addWidget(self.config_group)

    def _create_orchestrator_config(self):
        self.orch_config_widget = QWidget()
        layout = QVBoxLayout(self.orch_config_widget)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(12)

        grid = QGridLayout()
        grid.setHorizontalSpacing(10)
        grid.setVerticalSpacing(12)

        # Row 0: Host & Port
        grid.addWidget(self._create_label("Host"), 0, 0)
        self.orch_host = QLineEdit("0.0.0.0")
        self.orch_host.setFixedWidth(120)
        grid.addWidget(self.orch_host, 0, 1)
        grid.addWidget(self._create_label("Port"), 0, 2)
        self.orch_port = QSpinBox()
        self.orch_port.setRange(1, 65535)
        self.orch_port.setValue(DEFAULT_ORCH_PORT)
        self.orch_port.setFixedWidth(120)
        grid.addWidget(self.orch_port, 0, 3)

        # Row 1: Nodes & Cut Layer & Epochs
        grid.addWidget(self._create_label("Nodes"), 1, 0)
        self.n_nodes = QSpinBox()
        self.n_nodes.setRange(1, 10)
        self.n_nodes.setValue(2)
        self.n_nodes.setFixedWidth(120)
        grid.addWidget(self.n_nodes, 1, 1)
        grid.addWidget(self._create_label("Cut Layer"), 1, 2)
        self.cut_layer = QSpinBox()
        self.cut_layer.setRange(1, 20)
        self.cut_layer.setValue(5)
        self.cut_layer.setFixedWidth(120)
        grid.addWidget(self.cut_layer, 1, 3)
        grid.addWidget(self._create_label("Epochs"), 1, 4)
        self.epochs = QSpinBox()
        self.epochs.setRange(1, 1000)
        self.epochs.setValue(200)
        self.epochs.setFixedWidth(120)
        grid.addWidget(self.epochs, 1, 5)

        # Row 2: LR & Batch Size & Patience
        grid.addWidget(self._create_label("Learning Rate"), 2, 0)
        self.lr = QDoubleSpinBox()
        self.lr.setRange(0.0001, 1.0)
        self.lr.setDecimals(4)
        self.lr.setSingleStep(0.01)
        self.lr.setValue(0.1)
        self.lr.setFixedWidth(120)
        grid.addWidget(self.lr, 2, 1)
        grid.addWidget(self._create_label("Batch Size"), 2, 2)
        self.batch_size = QSpinBox()
        self.batch_size.setRange(1, 512)
        self.batch_size.setValue(128)
        self.batch_size.setFixedWidth(120)
        grid.addWidget(self.batch_size, 2, 3)
        grid.addWidget(self._create_label("Patience"), 2, 4)
        self.patience = QSpinBox()
        self.patience.setRange(1, 100)
        self.patience.setValue(30)
        self.patience.setFixedWidth(120)
        grid.addWidget(self.patience, 2, 5)

        # Row 3: Helper Port
        self.orch_helper_port_label = self._create_label("Helper Port")
        grid.addWidget(self.orch_helper_port_label, 3, 0)
        self.orch_helper_port = QSpinBox()
        self.orch_helper_port.setRange(1, 65535)
        self.orch_helper_port.setValue(DEFAULT_HELPER_ORCH_PORT)
        self.orch_helper_port.setFixedWidth(120)
        grid.addWidget(self.orch_helper_port, 3, 1)

        # Row 4: Checkboxes
        self.orch_secure = QCheckBox("Enable Secure Mode (MPC)")
        self.orch_secure.setChecked(True)
        self.orch_secure.toggled.connect(self._toggle_orch_secure)
        grid.addWidget(self.orch_secure, 4, 1, 1, 2)
        self.orch_no_accel = QCheckBox("CPU Only (disable GPU acceleration)")
        grid.addWidget(self.orch_no_accel, 4, 3, 1, 2)

        layout.addLayout(grid)

        self.config_layout.addWidget(self.orch_config_widget)

    def _create_label(self, text):
        label = QLabel(text)
        label.setStyleSheet(f"color: {COLORS['text_secondary']}; font-weight: 500;")
        label.setAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter)
        return label

    def _toggle_orch_secure(self, checked):
        self.orch_helper_port.setEnabled(checked)
        opacity = "1.0" if checked else "0.4"
        self.orch_helper_port_label.setStyleSheet(f"color: {COLORS['text_secondary']}; opacity: {opacity};")

    def _create_helper_config(self):
        self.helper_config_widget = QWidget()
        layout = QVBoxLayout(self.helper_config_widget)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(12)

        grid = QGridLayout()
        grid.setHorizontalSpacing(10)
        grid.setVerticalSpacing(12)

        # Row 0: Node Listen Host & Port
        grid.addWidget(self._create_label("Node Listen Host"), 0, 0)
        self.helper_node_host = QLineEdit("0.0.0.0")
        self.helper_node_host.setFixedWidth(180)
        grid.addWidget(self.helper_node_host, 0, 1)
        grid.addWidget(self._create_label("Node Listen Port"), 0, 2)
        self.helper_node_port = QSpinBox()
        self.helper_node_port.setRange(1, 65535)
        self.helper_node_port.setValue(DEFAULT_HELPER_NODE_PORT)
        self.helper_node_port.setFixedWidth(120)
        grid.addWidget(self.helper_node_port, 0, 3)

        # Row 1: Orchestrator Host & Port
        grid.addWidget(self._create_label("Orchestrator Host"), 1, 0)
        self.helper_orch_host = QLineEdit("127.0.0.1")
        self.helper_orch_host.setFixedWidth(180)
        grid.addWidget(self.helper_orch_host, 1, 1)
        grid.addWidget(self._create_label("Orchestrator Port"), 1, 2)
        self.helper_orch_port = QSpinBox()
        self.helper_orch_port.setRange(1, 65535)
        self.helper_orch_port.setValue(DEFAULT_HELPER_ORCH_PORT)
        self.helper_orch_port.setFixedWidth(120)
        grid.addWidget(self.helper_orch_port, 1, 3)

        # Row 2: Number of Nodes
        grid.addWidget(self._create_label("Number of Nodes"), 2, 0)
        self.helper_n_nodes = QSpinBox()
        self.helper_n_nodes.setRange(1, 10)
        self.helper_n_nodes.setValue(2)
        self.helper_n_nodes.setFixedWidth(120)
        grid.addWidget(self.helper_n_nodes, 2, 1)

        center = QHBoxLayout()
        center.addStretch()
        center.addLayout(grid)
        center.addStretch()
        layout.addLayout(center)

        self.config_layout.addWidget(self.helper_config_widget)

    def _create_node_config(self):
        self.node_config_widget = QWidget()
        layout = QVBoxLayout(self.node_config_widget)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(12)

        grid = QGridLayout()
        grid.setHorizontalSpacing(10)
        grid.setVerticalSpacing(12)

        # Row 0: Orchestrator Host & Port
        grid.addWidget(self._create_label("Orchestrator Host"), 0, 0)
        self.node_orch_host = QLineEdit("127.0.0.1")
        self.node_orch_host.setFixedWidth(180)
        grid.addWidget(self.node_orch_host, 0, 1)
        grid.addWidget(self._create_label("Orchestrator Port"), 0, 2)
        self.node_orch_port = QSpinBox()
        self.node_orch_port.setRange(1, 65535)
        self.node_orch_port.setValue(DEFAULT_ORCH_PORT)
        self.node_orch_port.setFixedWidth(120)
        grid.addWidget(self.node_orch_port, 0, 3)

        # Row 1: Helper Host & Port
        self.node_helper_host_label = self._create_label("Helper Host")
        grid.addWidget(self.node_helper_host_label, 1, 0)
        self.node_helper_host = QLineEdit("127.0.0.1")
        self.node_helper_host.setFixedWidth(180)
        grid.addWidget(self.node_helper_host, 1, 1)
        self.node_helper_port_label = self._create_label("Helper Port")
        grid.addWidget(self.node_helper_port_label, 1, 2)
        self.node_helper_port = QSpinBox()
        self.node_helper_port.setRange(1, 65535)
        self.node_helper_port.setValue(DEFAULT_HELPER_NODE_PORT)
        self.node_helper_port.setFixedWidth(120)
        grid.addWidget(self.node_helper_port, 1, 3)

        # Row 2: Node ID
        grid.addWidget(self._create_label("Node ID"), 2, 0)
        self.node_id_label = QLabel("Auto-assigned by Orchestrator")
        self.node_id_label.setStyleSheet(f"color: {COLORS['text_muted']}; font-style: italic;")
        grid.addWidget(self.node_id_label, 2, 1, 1, 3)

        center = QHBoxLayout()
        center.addStretch()
        center.addLayout(grid)
        center.addStretch()
        layout.addLayout(center)

        # Checkboxes - grouped in a single VBox then centered
        cb_box = QVBoxLayout()
        cb_box.setSpacing(8)

        self.node_secure = QCheckBox("Enable Secure Mode (MPC)")
        self.node_secure.setChecked(True)
        self.node_secure.toggled.connect(self._toggle_node_secure)
        cb_box.addWidget(self.node_secure)

        cb_row2 = QHBoxLayout()
        self.node_isolated = QCheckBox("Isolated Mode (process isolation)")
        self.node_isolated.setToolTip("Run with process isolation for enhanced security")
        cb_row2.addWidget(self.node_isolated)
        cb_row2.addSpacing(30)
        self.node_no_accel = QCheckBox("CPU Only (disable GPU acceleration)")
        cb_row2.addWidget(self.node_no_accel)
        cb_row2.addStretch()
        cb_box.addLayout(cb_row2)

        cb_center = QHBoxLayout()
        cb_center.addStretch()
        cb_center.addLayout(cb_box)
        cb_center.addStretch()
        layout.addLayout(cb_center)

        self.config_layout.addWidget(self.node_config_widget)

    def _toggle_node_secure(self, checked):
        self.node_helper_host.setEnabled(checked)
        self.node_helper_port.setEnabled(checked)

    def _on_role_change(self, role):
        self.current_role = role

        self.orch_config_widget.hide()
        self.helper_config_widget.hide()
        self.node_config_widget.hide()

        # Training progress dashboard is Orchestrator-only
        self.progress_group.setVisible(role == 'orchestrator')

        if role == 'orchestrator':
            self.orch_config_widget.show()
        elif role == 'helper':
            self.helper_config_widget.show()
        else:
            self.node_config_widget.show()

    def _create_status(self, parent_layout):
        frame = QFrame()
        frame.setObjectName("statusFrame")
        layout = QHBoxLayout(frame)
        layout.setContentsMargins(15, 10, 15, 10)

        # Status indicator
        self.status_dot = QLabel("●")
        self.status_dot.setStyleSheet(f"font-size: 16px; color: {COLORS['text_muted']};")
        layout.addWidget(self.status_dot)

        self.status_label = QLabel("Ready")
        self.status_label.setStyleSheet(f"font-size: 14px; font-weight: bold; color: {COLORS['text_primary']};")
        layout.addWidget(self.status_label)

        self.detail_label = QLabel("")
        self.detail_label.setStyleSheet(f"font-size: 12px; color: {COLORS['text_muted']};")
        layout.addWidget(self.detail_label)

        layout.addStretch()

        # 실행 모드 표시
        mode_text = "🐳 Docker" if self.use_docker else "🖥️ Native"
        self.mode_label = QLabel(mode_text)
        self.mode_label.setStyleSheet(f"font-size: 12px; color: {COLORS['text_muted']}; margin-right: 15px;")
        layout.addWidget(self.mode_label)

        # IP display
        ips = get_local_ip_addresses()
        ip_label = QLabel(f"📡 Your IP: {', '.join(ips)}")
        ip_label.setStyleSheet(f"font-size: 12px; color: {COLORS['accent_info']};")
        layout.addWidget(ip_label)

        parent_layout.addWidget(frame)

    def _create_controls(self, parent_layout):
        layout = QHBoxLayout()
        layout.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.setSpacing(20)

        self.start_btn = QPushButton("▶  Start Training")
        self.start_btn.setFixedSize(180, 45)
        self.start_btn.clicked.connect(self._start_training)
        layout.addWidget(self.start_btn)

        self.stop_btn = QPushButton("■  Stop")
        self.stop_btn.setObjectName("stopBtn")
        self.stop_btn.setFixedSize(140, 45)
        self.stop_btn.setEnabled(False)
        self.stop_btn.clicked.connect(self._stop_training)
        layout.addWidget(self.stop_btn)

        parent_layout.addLayout(layout)

    def _create_training_progress(self, parent_layout):
        self.progress_group = QGroupBox("Training Progress")
        layout = QVBoxLayout()
        layout.setSpacing(12)

        # Row 1: Progress bars (stacked) + elapsed time
        top_row = QHBoxLayout()

        bars_col = QVBoxLayout()
        bars_col.setSpacing(0)

        self.epoch_progress = QProgressBar()
        self.epoch_progress.setRange(0, 100)
        self.epoch_progress.setValue(0)
        self.epoch_progress.setFormat("%v / %m")
        bars_col.addWidget(self.epoch_progress)

        self.batch_progress = QProgressBar()
        self.batch_progress.setObjectName("batchBar")
        self.batch_progress.setRange(0, 100)
        self.batch_progress.setValue(0)
        self.batch_progress.setTextVisible(False)
        bars_col.addWidget(self.batch_progress)

        top_row.addLayout(bars_col, stretch=1)

        top_row.addSpacing(15)

        self.elapsed_label = QLabel("00:00:00")
        self.elapsed_label.setStyleSheet(
            f"font-size: 14px; font-weight: bold; "
            f"color: {COLORS['text_secondary']}; font-family: monospace;"
        )
        top_row.addWidget(self.elapsed_label)

        layout.addLayout(top_row)

        # Row 2-3: Metric cards grid (2 rows x 3 columns)
        grid = QGridLayout()
        grid.setSpacing(10)

        card_defs = [
            ('epoch', 'Epoch', '--', None),
            ('train_loss', 'Train Loss', '--', COLORS['accent_warning']),
            ('test_loss', 'Test Loss', '--', COLORS['accent_info']),
            ('test_acc', 'Test Accuracy', '--%', COLORS['accent_success']),
            ('best_acc', 'Best Accuracy', '--%', COLORS['accent_primary']),
            ('lr', 'Learning Rate', '--', COLORS['text_secondary']),
        ]

        for idx, (key, title, init_val, color) in enumerate(card_defs):
            card = MetricCard(title, init_val, accent_color=color)
            self._metric_cards[key] = card
            grid.addWidget(card, idx // 3, idx % 3)

        layout.addLayout(grid)

        self.progress_group.setLayout(layout)
        parent_layout.addWidget(self.progress_group)

    def _create_log_display(self, parent_layout):
        group = QGroupBox("Training Logs")
        layout = QVBoxLayout()

        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setMinimumHeight(180)
        layout.addWidget(self.log_text)

        # Controls
        ctrl_layout = QHBoxLayout()

        clear_btn = QPushButton("Clear")
        clear_btn.setObjectName("clearBtn")
        clear_btn.clicked.connect(self.log_text.clear)
        ctrl_layout.addWidget(clear_btn)

        ctrl_layout.addStretch()

        self.auto_scroll = QCheckBox("Auto-scroll")
        self.auto_scroll.setChecked(True)
        ctrl_layout.addWidget(self.auto_scroll)

        layout.addLayout(ctrl_layout)

        group.setLayout(layout)
        parent_layout.addWidget(group)

    def _get_config(self):
        """Build config dict based on current role."""
        if self.current_role == 'orchestrator':
            return {
                'host': self.orch_host.text(),
                'port': self.orch_port.value(),
                'n_nodes': self.n_nodes.value(),
                'cut_layer': self.cut_layer.value(),
                'epochs': self.epochs.value(),
                'lr': self.lr.value(),
                'train_batch_size': self.batch_size.value(),
                'test_batch_size': 256,
                'patience': self.patience.value(),
                'secure': self.orch_secure.isChecked(),
                'helper_host': self.orch_host.text(),  # Host와 동일하게 자동 설정
                'helper_port': self.orch_helper_port.value(),
                'momentum': 0.9,
                'weight_decay': 5e-4,
                'nesterov': True,
                'scheduler': 'cosine',
                'eta_min': 1e-6,
                'grad_clip_norm': 1.0,
                'no_accel': self.orch_no_accel.isChecked(),
            }

        elif self.current_role == 'helper':
            return {
                'node_host': self.helper_node_host.text(),
                'node_port': self.helper_node_port.value(),
                'orch_host': self.helper_orch_host.text(),
                'orch_port': self.helper_orch_port.value(),
                'n_nodes': self.helper_n_nodes.value(),
            }

        else:  # node
            secure = self.node_secure.isChecked()
            return {
                'host': self.node_orch_host.text(),
                'port': self.node_orch_port.value(),
                'secure': secure,
                'orch_host': self.node_orch_host.text(),
                'orch_port': self.node_orch_port.value(),
                'helper_host': self.node_helper_host.text() if secure else '127.0.0.1',
                'helper_port': self.node_helper_port.value() if secure else 8081,
                'no_accel': self.node_no_accel.isChecked(),
                'isolated': self.node_isolated.isChecked(),
            }

    def _validate_config(self, config):
        """Validate configuration."""
        errors = []

        if self.current_role == 'orchestrator':
            if config.get('n_nodes', 0) < 1:
                errors.append("Number of nodes must be at least 1")
            if config.get('epochs', 0) < 1:
                errors.append("Epochs must be at least 1")
            if config.get('lr', 0) <= 0:
                errors.append("Learning rate must be positive")

        elif self.current_role == 'helper':
            if config.get('n_nodes', 0) < 1:
                errors.append("Number of nodes must be at least 1")

        if errors:
            QMessageBox.critical(self, "Configuration Error", "\n".join(errors))
            return False

        return True

    def _check_docker_on_startup(self):
        """시작 시 Docker 상태 확인"""
        if not self.docker_manager:
            return

        if not self.docker_manager.is_docker_installed():
            # Docker 미설치 - 상태 표시만 하고 넘어감
            self._update_status('idle', 'Docker 설치 필요')
            return

        if not self.docker_manager.is_docker_running():
            self._update_status('idle', 'Docker 시작 필요')
            return

        self.docker_ready = True
        if self.docker_manager.is_image_available():
            self._update_status('idle', 'Ready (Docker)')
        else:
            self._update_status('idle', 'Ready (이미지 빌드 필요)')

    def _ensure_docker_ready(self) -> bool:
        """
        Docker 환경이 준비되었는지 확인하고 필요시 다이얼로그 표시.
        Returns: True if ready, False otherwise
        """
        if not self.use_docker:
            return True  # 네이티브 모드

        if not self.docker_manager:
            QMessageBox.critical(
                self, "오류",
                "Docker 관리 모듈을 불러올 수 없습니다.\n"
                "docker_manager.py 파일이 있는지 확인해주세요."
            )
            return False

        # Docker 설치 확인
        while not self.docker_manager.is_docker_installed():
            dialog = DockerSetupDialog(
                self,
                download_url=self.docker_manager.get_docker_download_url()
            )
            if dialog.exec() != QDialog.DialogCode.Accepted:
                return False  # 사용자가 취소함

        # Docker 실행 확인
        while not self.docker_manager.is_docker_running():
            dialog = DockerNotRunningDialog(self)
            if dialog.exec() != QDialog.DialogCode.Accepted:
                return False  # 사용자가 취소함

        # 이미지 확인 및 빌드
        if not self.docker_manager.is_image_available():
            dialog = ImageBuildDialog(self.docker_manager, self)
            dialog.start_build()
            dialog.exec()

            if not dialog.build_success:
                QMessageBox.critical(
                    self, "빌드 실패",
                    "Docker 이미지 빌드에 실패했습니다.\n"
                    "인터넷 연결을 확인하고 다시 시도해주세요."
                )
                return False

        self.docker_ready = True
        return True

    def _start_training(self):
        if self.is_running:
            return

        config = self._get_config()

        if not self._validate_config(config):
            return

        # Auto-save current config for next launch
        try:
            self._save_config_to_file(LAST_CONFIG_PATH)
        except Exception:
            pass

        # Clear logs and reset metrics
        self.log_text.clear()
        if self.current_role == 'orchestrator':
            self._reset_metrics()
            self._training_start_time = time.time()
            self.elapsed_timer.start()

        # Reset node ID display
        if self.current_role == 'node':
            self.node_id_label.setText("Connecting...")
            self.node_id_label.setStyleSheet(f"color: {COLORS['accent_warning']};")

        # Create and start training thread
        self.training_thread = TrainingThread(
            role=self.current_role,
            config=config,
            message_queue=self.message_queue,
            use_docker=self.use_docker,
            docker_manager=self.docker_manager
        )
        self.training_thread.start()

        self.is_running = True
        self._set_running_state(True)
        mode_text = "(Docker)" if self.use_docker else "(Native)"
        self._update_status('setup', f'Starting {mode_text}...')

    def _stop_training(self):
        if not self.is_running:
            return

        self._update_status('stopping', '')

        if self.training_thread:
            self.training_thread.request_stop()

    def _set_running_state(self, running):
        self.start_btn.setEnabled(not running)
        self.stop_btn.setEnabled(running)

        # Disable role selection and config when running
        for card in self.role_cards.values():
            card.setEnabled(not running)

        self.orch_config_widget.setEnabled(not running)
        self.helper_config_widget.setEnabled(not running)
        self.node_config_widget.setEnabled(not running)

    def _update_status(self, state, detail):
        states = {
            'idle': ('Ready', COLORS['text_muted']),
            'setup': ('Setting up...', COLORS['accent_warning']),
            'waiting': ('Waiting...', COLORS['accent_warning']),
            'running': ('Running', COLORS['accent_success']),
            'stopping': ('Stopping...', COLORS['accent_warning']),
            'stopped': ('Stopped', COLORS['text_muted']),
            'completed': ('Completed', COLORS['accent_info']),
            'error': ('Error', COLORS['accent_primary']),
        }

        text, color = states.get(state, (state, COLORS['text_muted']))
        self.status_label.setText(text)
        self.status_label.setStyleSheet(f"font-size: 14px; font-weight: bold; color: {color};")
        self.status_dot.setStyleSheet(f"font-size: 16px; color: {color};")
        self.detail_label.setText(detail or "")

    def _poll_queue(self):
        """Poll message queue for updates."""
        try:
            while True:
                msg = self.message_queue.get_nowait()
                self._handle_message(msg)
        except queue.Empty:
            pass

    def _handle_message(self, msg):
        msg_type = msg[0]

        if msg_type == 'log':
            _, level, text = msg
            self._append_log(text, level)
            if self.current_role == 'orchestrator':
                self._parse_metrics(text)

        elif msg_type == 'status':
            _, state, detail = msg
            self._update_status(state, detail)

            if state in ('stopped', 'completed', 'error'):
                self.is_running = False
                self._set_running_state(False)
                self.elapsed_timer.stop()

        elif msg_type == 'node_id':
            _, node_id, _ = msg
            self.node_id_label.setText(f"Node {node_id}")
            self.node_id_label.setStyleSheet(f"color: {COLORS['accent_success']}; font-weight: bold;")

    # ------------------------------------------------------------------
    # Training metrics
    # ------------------------------------------------------------------

    _RE_EPOCH = re.compile(r'Epoch\s+(\d+)/(\d+)')
    _RE_LR = re.compile(r'LR:\s*([\d.]+)')
    _RE_TRAIN_LOSS = re.compile(r'Train Loss:\s*([\d.]+)')
    _RE_TEST_LOSS = re.compile(r'Test Loss:\s*([\d.]+)')
    _RE_TEST_ACC = re.compile(r'Test Acc:\s*([\d.]+)%')
    _RE_BEST_ACC = re.compile(r'best accuracy:\s*([\d.]+)%')
    _RE_PATIENCE = re.compile(r'No improvement \((\d+)/(\d+)\)')
    _RE_BATCH = re.compile(r'(\d+)%\|[^|]*\|\s*(\d+)/(\d+)')

    def _parse_metrics(self, text):
        """Extract training metrics from a log line and update dashboard."""
        # Batch progress from tqdm output (minute hand)
        m = self._RE_BATCH.search(text)
        if m:
            cur_batch, total_batch = int(m.group(2)), int(m.group(3))
            self.batch_progress.setRange(0, total_batch)
            self.batch_progress.setValue(cur_batch)
            return  # tqdm lines don't contain other metrics

        # Epoch progress (hour hand) — resets batch bar on new epoch
        m = self._RE_EPOCH.search(text)
        if m:
            cur, total = int(m.group(1)), int(m.group(2))
            self.epoch_progress.setRange(0, total)
            self.epoch_progress.setValue(cur)
            self.epoch_progress.setFormat(f"{cur} / {total}")
            self._metric_cards['epoch'].set_value(f"{cur} / {total}")
            self.batch_progress.setValue(0)

        m = self._RE_LR.search(text)
        if m:
            self._metric_cards['lr'].set_value(m.group(1))

        m = self._RE_TRAIN_LOSS.search(text)
        if m:
            self._metric_cards['train_loss'].set_value(m.group(1))

        m = self._RE_TEST_LOSS.search(text)
        if m:
            self._metric_cards['test_loss'].set_value(m.group(1))

        m = self._RE_TEST_ACC.search(text)
        if m:
            self._metric_cards['test_acc'].set_value(f"{m.group(1)}%")

        m = self._RE_BEST_ACC.search(text)
        if m:
            self._metric_cards['best_acc'].set_value(f"{m.group(1)}%")

    def _reset_metrics(self):
        """Reset all metric cards and progress bar to initial state."""
        self.epoch_progress.setRange(0, 100)
        self.epoch_progress.setValue(0)
        self.epoch_progress.setFormat("0 / --")
        self.batch_progress.setRange(0, 100)
        self.batch_progress.setValue(0)
        self.elapsed_label.setText("00:00:00")

        defaults = {
            'epoch': '--',
            'train_loss': '--',
            'test_loss': '--',
            'test_acc': '--%',
            'best_acc': '--%',
            'lr': '--',
        }
        for key, val in defaults.items():
            self._metric_cards[key].set_value(val)

    def _update_elapsed_time(self):
        """Update elapsed time label every second."""
        if self._training_start_time is None:
            return
        elapsed = int(time.time() - self._training_start_time)
        h, remainder = divmod(elapsed, 3600)
        m, s = divmod(remainder, 60)
        self.elapsed_label.setText(f"{h:02d}:{m:02d}:{s:02d}")

    # ------------------------------------------------------------------
    # Config save / load
    # ------------------------------------------------------------------

    def _collect_config_for_save(self):
        """Collect all widget values into a JSON-serialisable dict."""
        data = {'role': self.current_role}
        if self.current_role == 'orchestrator':
            data.update({
                'host': self.orch_host.text(),
                'port': self.orch_port.value(),
                'n_nodes': self.n_nodes.value(),
                'cut_layer': self.cut_layer.value(),
                'epochs': self.epochs.value(),
                'lr': self.lr.value(),
                'batch_size': self.batch_size.value(),
                'patience': self.patience.value(),
                'secure': self.orch_secure.isChecked(),
                'helper_port': self.orch_helper_port.value(),
                'no_accel': self.orch_no_accel.isChecked(),
            })
        elif self.current_role == 'helper':
            data.update({
                'node_host': self.helper_node_host.text(),
                'node_port': self.helper_node_port.value(),
                'orch_host': self.helper_orch_host.text(),
                'orch_port': self.helper_orch_port.value(),
                'n_nodes': self.helper_n_nodes.value(),
            })
        else:
            data.update({
                'orch_host': self.node_orch_host.text(),
                'orch_port': self.node_orch_port.value(),
                'secure': self.node_secure.isChecked(),
                'helper_host': self.node_helper_host.text(),
                'helper_port': self.node_helper_port.value(),
                'no_accel': self.node_no_accel.isChecked(),
                'isolated': self.node_isolated.isChecked(),
            })
        return data

    def _apply_config_from_dict(self, data):
        """Populate widgets from a previously saved config dict."""
        role = data.get('role', self.current_role)
        self.select_role(role)

        if role == 'orchestrator':
            self.orch_host.setText(data.get('host', '0.0.0.0'))
            self.orch_port.setValue(data.get('port', DEFAULT_ORCH_PORT))
            self.n_nodes.setValue(data.get('n_nodes', 2))
            self.cut_layer.setValue(data.get('cut_layer', 5))
            self.epochs.setValue(data.get('epochs', 200))
            self.lr.setValue(data.get('lr', 0.1))
            self.batch_size.setValue(data.get('batch_size', 128))
            self.patience.setValue(data.get('patience', 30))
            self.orch_secure.setChecked(data.get('secure', True))
            self.orch_helper_port.setValue(data.get('helper_port', DEFAULT_HELPER_ORCH_PORT))
            self.orch_no_accel.setChecked(data.get('no_accel', False))
        elif role == 'helper':
            self.helper_node_host.setText(data.get('node_host', '0.0.0.0'))
            self.helper_node_port.setValue(data.get('node_port', DEFAULT_HELPER_NODE_PORT))
            self.helper_orch_host.setText(data.get('orch_host', '127.0.0.1'))
            self.helper_orch_port.setValue(data.get('orch_port', DEFAULT_HELPER_ORCH_PORT))
            self.helper_n_nodes.setValue(data.get('n_nodes', 2))
        else:
            self.node_orch_host.setText(data.get('orch_host', '127.0.0.1'))
            self.node_orch_port.setValue(data.get('orch_port', DEFAULT_ORCH_PORT))
            self.node_secure.setChecked(data.get('secure', True))
            self.node_helper_host.setText(data.get('helper_host', '127.0.0.1'))
            self.node_helper_port.setValue(data.get('helper_port', DEFAULT_HELPER_NODE_PORT))
            self.node_no_accel.setChecked(data.get('no_accel', False))
            self.node_isolated.setChecked(data.get('isolated', False))

    def _save_config_to_file(self, path):
        """Write current config to a JSON file."""
        data = self._collect_config_for_save()
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _load_config_from_file(self, path):
        """Read a JSON config file and apply it to widgets."""
        with open(path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        self._apply_config_from_dict(data)

    def _on_save_preset(self):
        path, _ = QFileDialog.getSaveFileName(
            self, "Save Preset", str(CONFIG_DIR),
            "JSON Files (*.json);;All Files (*)"
        )
        if path:
            try:
                self._save_config_to_file(path)
            except Exception as e:
                QMessageBox.warning(self, "Save Error", str(e))

    def _on_load_preset(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "Load Preset", str(CONFIG_DIR),
            "JSON Files (*.json);;All Files (*)"
        )
        if path:
            try:
                self._load_config_from_file(path)
            except Exception as e:
                QMessageBox.warning(self, "Load Error", str(e))

    def _restore_last_config(self):
        """Silently restore the last-used config on app launch."""
        if LAST_CONFIG_PATH.exists():
            try:
                self._load_config_from_file(LAST_CONFIG_PATH)
            except Exception:
                pass

    def _append_log(self, text, level):
        colors = {
            'INFO': COLORS['text_primary'],
            'WARNING': COLORS['accent_warning'],
            'ERROR': COLORS['accent_primary'],
            'DEBUG': COLORS['text_muted'],
        }
        color = colors.get(level, COLORS['text_primary'])

        self.log_text.setTextColor(QColor(color))
        self.log_text.append(text)

        if self.auto_scroll.isChecked():
            self.log_text.moveCursor(QTextCursor.MoveOperation.End)

    def closeEvent(self, event):
        if self.is_running:
            reply = QMessageBox.question(
                self, 'Confirm Exit',
                'Training is in progress. Stop and exit?',
                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
            )
            if reply == QMessageBox.StandardButton.Yes:
                self._stop_training()
                # Docker 정리
                if self.docker_manager:
                    self.docker_manager.cleanup()
                event.accept()
            else:
                event.ignore()
        else:
            # Docker 정리
            if self.docker_manager:
                self.docker_manager.cleanup()
            event.accept()

    def _toggle_theme(self):
        """Toggle between dark and light themes."""
        global COLORS

        self.is_dark_mode = not self.is_dark_mode

        if self.is_dark_mode:
            COLORS = COLORS_DARK.copy()
            self.theme_btn.setText("☀️")
            self.theme_btn.setToolTip("Switch to Light Mode")
        else:
            COLORS = COLORS_LIGHT.copy()
            self.theme_btn.setText("🌙")
            self.theme_btn.setToolTip("Switch to Dark Mode")

        # Apply new stylesheet to application
        QApplication.instance().setStyleSheet(get_stylesheet(COLORS))

        # Update main backgrounds
        self.scroll_area.setStyleSheet(f"""
            QScrollArea {{
                background-color: {COLORS['bg_dark']};
                border: none;
            }}
        """)
        self.content_widget.setStyleSheet(f"background-color: {COLORS['bg_dark']};")

        # Update header styles
        self.title_header.setStyleSheet(f"""
            font-size: 36px;
            font-weight: bold;
            color: {COLORS['accent_primary']};
        """)
        self.subtitle_header.setStyleSheet(f"""
            font-size: 14px;
            color: {COLORS['text_muted']};
            margin-top: -1px;
        """)

        # Update role cards
        for card in self.role_cards.values():
            card._update_labels()
            card.set_selected(card == self.role_cards.get(self.current_role))

        # Update status indicators
        self.status_dot.setStyleSheet(f"font-size: 16px; color: {COLORS['text_muted']};")
        self.status_label.setStyleSheet(f"font-size: 14px; font-weight: bold; color: {COLORS['text_primary']};")
        self.detail_label.setStyleSheet(f"font-size: 12px; color: {COLORS['text_muted']};")

        # Update metric cards
        for card in self._metric_cards.values():
            card.update_style()
        self.elapsed_label.setStyleSheet(
            f"font-size: 14px; font-weight: bold; "
            f"color: {COLORS['text_secondary']}; font-family: monospace;"
        )

        # Update node ID label if in node mode
        if self.current_role == 'node':
            self.node_id_label.setStyleSheet(f"color: {COLORS['text_muted']}; font-style: italic;")


# ==============================================================================
# ENTRY POINT
# ==============================================================================

def main():
    app = QApplication(sys.argv)
    app.setStyleSheet(STYLESHEET)

    window = TLPPMainWindow()
    window.show()

    sys.exit(app.exec())


if __name__ == '__main__':
    main()