main.py

import os
import sys
import subprocess
import shutil
import argparse
import logging
from monitor.utils import get_logger

logger = get_logger(__name__)

VENV_DIR = "orch_monitoring_venv"
REQUIREMENTS = "requirements.txt"

def create_venv():
    """
    Create a Python virtual environment.

    Exits if Python is not installed.

    Returns:
        None
    """
    python_cmd = "python3" if os.name != "nt" else "python"
    if shutil.which(python_cmd) is None:
        logger.error(f"{python_cmd} is not installed.")
        sys.exit(1)
    subprocess.check_call([sys.executable, "-m", "venv", VENV_DIR])

def install_requirements():
    """
    Install required Python packages into the virtual environment.

    Returns:
        None
    """
    pip_path = os.path.join(VENV_DIR, "bin", "pip") if os.name != "nt" else os.path.join(VENV_DIR, "Scripts", "pip.exe")
    subprocess.check_call([pip_path, "install", "-r", REQUIREMENTS])

def run_script_in_venv(logfile=None):
    """
    Re-run the current script inside the virtual environment.

    Args:
        logfile (str, optional): Path to log file for output. If None, logs to console only.

    Returns:
        None
    """
    python_path = os.path.join(VENV_DIR, "bin", "python") if os.name != "nt" else os.path.join(VENV_DIR, "Scripts", "python.exe")
    cmd = [python_path, __file__]
    if logfile:
        cmd += ["--logfile", logfile]
    subprocess.check_call(cmd)

def main():
    """
    Main entry point for the monitoring tool.

    Loads configuration and starts the appropriate monitor based on config.

    Returns:
        None
    """
    from monitor.utils import load_config
    import monitor.gms_monitor as Gmonitor
    import monitor.edge_connect_monitor as ECmonitor

    config = load_config()

    if config['ORCHESTRATOR'].getboolean('GMS_monitor'):
        monitor = Gmonitor.GMSMonitor(config)
        monitor.run()
    elif config['EDGE_CONNECT'].getboolean('EC_monitor'):
        monitor = ECmonitor.ECMonitor(config)
        monitor.run()
    else:
        logger.error("Please set 'GMS_monitor' or 'EC_monitor' to 'True' in the 'config.ini' file.")
        sys.exit(1)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Orchestrator Monitoring Tool",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""
Examples:
  python main.py                # Run with default console logging
  python main.py --logfile output.log   # Log all output to output.log as well as console
  python main.py --help         # Show this help message
        """
    )
    parser.add_argument(
        "--logfile",
        type=str,
        help="Path to log file for all terminal output. If not provided, logs are printed to console only."
    )
    args, unknown = parser.parse_known_args()

    # Set up logging and redirect stdout/stderr to logfile if specified
    if args.logfile and args.logfile.strip():
        log_dir = os.path.dirname(os.path.abspath(args.logfile))
        if log_dir and not os.path.exists(log_dir):
            os.makedirs(log_dir)
        # Set up logging handlers for file and console
        root_logger = logging.getLogger()
        for handler in root_logger.handlers[:]:
            root_logger.removeHandler(handler)
        file_handler = logging.FileHandler(args.logfile, mode='w')
        stream_handler = logging.StreamHandler(sys.stdout)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        stream_handler.setFormatter(formatter)
        root_logger.addHandler(file_handler)
        root_logger.addHandler(stream_handler)
        root_logger.setLevel(logging.INFO)
        # Redirect stdout and stderr to logfile as well as console
        class TeeStream:
            """
            Stream that writes to multiple streams simultaneously.
            """
            def __init__(self, *streams):
                self.streams = streams
            def write(self, data):
                for s in self.streams:
                    s.write(data)
                    s.flush()
            def flush(self):
                for s in self.streams:
                    s.flush()
        logfile_stream = open(args.logfile, 'w', buffering=1)
        sys.stdout = TeeStream(sys.__stdout__, logfile_stream)
        sys.stderr = TeeStream(sys.__stderr__, logfile_stream)
    else:
        # No logfile argument provided, just use default logging to console
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    logger = get_logger(__name__)

    venv_python = os.path.abspath(os.path.join(VENV_DIR, "bin", "python")) if os.name != "nt" else os.path.abspath(os.path.join(VENV_DIR, "Scripts", "python.exe"))
    if not os.path.isdir(VENV_DIR):
        logger.info("Creating virtual environment...")
        create_venv()
        logger.info("Installing requirements...")
        install_requirements()
    if os.path.abspath(sys.executable) != venv_python:
        logger.info("Re-running script inside virtual environment...")
        try:
            run_script_in_venv(args.logfile)
        except KeyboardInterrupt:
            sys.exit(0)
    else:
        logger.info("Running inside virtual environment.")
        try:
            main()
        except KeyboardInterrupt:
            logger.info("Monitoring stopped by user.")
            sys.exit(0)