main.py

from __future__ import annotations

import logging
import os
import warnings
import sys
import signal


class _NullWriter:
    def write(self, *args, **kwargs):
        pass

    def flush(self, *args, **kwargs):
        pass


if sys.stdout is None:
    sys.stdout = _NullWriter()
if sys.stderr is None:
    sys.stderr = _NullWriter()

os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ["HF_HUB_DISABLE_XET"] = "1"

warnings.filterwarnings("ignore", category=UserWarning, message=r".*pkg_resources is deprecated.*")
warnings.filterwarnings("ignore", message=r".*Megatron.*")
warnings.filterwarnings("ignore", message=r".*TensorFloat-32.*")
warnings.filterwarnings("ignore", message=r".*TF32.*")
warnings.filterwarnings("ignore", message=r".*allow_tf32.*")
warnings.filterwarnings("ignore", message=r".*Redirects are currently not supported.*")

for _name in ["nemo", "nemo.collections", "nemo.utils", "nemo.core",
              "lightning", "lightning_fabric", "pytorch_lightning",
              "nemo_logger", "nv_one_logger", "wandb", "numba", "httpx"]:
    logging.getLogger(_name).setLevel(logging.ERROR)

from core.cuda_setup import set_cuda_paths
_cuda_paths_configured = set_cuda_paths()

from PySide6.QtWidgets import QApplication, QMessageBox

from core.logging_config import setup_logging, get_logger
from core.temp_file_manager import temp_file_manager
from gui.main_window import MainWindow
from gui.styles import APP_STYLESHEET


def _install_sigint_handler() -> None:
    signal.signal(signal.SIGINT, signal.SIG_DFL)


def _flush_clipboard_before_hard_exit() -> None:
    if sys.platform == "win32":
        try:
            import ctypes
            ctypes.windll.ole32.OleFlushClipboard()
        except Exception:
            pass


def _global_exception_handler(exc_type, exc_value, exc_tb):
    logger = get_logger(__name__)
    logger.critical("Unhandled exception", exc_info=(exc_type, exc_value, exc_tb))

    app = QApplication.instance()
    if app:
        QMessageBox.critical(
            None,
            "Critical Error",
            f"An unexpected error occurred:\n\n{exc_value}\n\nPlease check the log file for details."
        )


def _check_cuda_available() -> bool:
    try:
        import torch
        return torch.cuda.is_available()
    except Exception:
        return False


def _get_cuda_device_name() -> str | None:
    try:
        import torch
        if torch.cuda.is_available():
            return torch.cuda.get_device_name(0)
    except Exception:
        pass
    return None


def _configure_torch_threads(logger) -> None:
    try:
        import psutil
        import torch
        physical = psutil.cpu_count(logical=False)
        if physical and physical > 0:
            n_threads = max(4, physical - 4)
            torch.set_num_threads(n_threads)
            logger.info(
                f"torch.set_num_threads({n_threads}) "
                f"(physical={physical}, logical={psutil.cpu_count(logical=True)}, "
                f"formula=max(4, physical-4))"
            )
        else:
            logger.warning("Could not detect physical core count; leaving torch threads at default")
    except Exception as e:
        logger.warning(f"Failed to configure torch threads: {e}")


def run_gui() -> None:
    log_file = setup_logging()
    logger = get_logger(__name__)
    logger.info("Application starting")
    logger.info(f"Log file: {log_file}")
    logger.info(f"CUDA paths configured: {_cuda_paths_configured}")

    _configure_torch_threads(logger)

    sys.excepthook = _global_exception_handler

    app = QApplication(sys.argv)
    app.setStyle('Fusion')
    app.setStyleSheet(APP_STYLESHEET)
    _install_sigint_handler()

    cuda_ok = _check_cuda_available()
    logger.info(f"CUDA available: {cuda_ok}")

    if cuda_ok:
        device_name = _get_cuda_device_name()
        if device_name:
            logger.info(f"CUDA device: {device_name}")

    try:
        window = MainWindow(cuda_available=cuda_ok)
        window.show()

        exit_code = app.exec()

        temp_file_manager.cleanup_all()
        logger.info(f"Application exiting with code {exit_code}")
        logging.shutdown()
        _flush_clipboard_before_hard_exit()
        os._exit(exit_code)

    except Exception as e:
        logger.critical(f"Failed to start application: {e}", exc_info=True)
        QMessageBox.critical(
            None,
            "Startup Error",
            f"Failed to start application:\n\n{e}"
        )
        temp_file_manager.cleanup_all()
        logging.shutdown()
        _flush_clipboard_before_hard_exit()
        os._exit(1)


def main() -> None:
    run_gui()


if __name__ == "__main__":
    main()