ETF_Suite_Portal/ETF_Portal/cli.py

#!/usr/bin/env python3
"""
ETF Portal CLI

A command-line interface for managing the ETF Portal application.
"""

import click
import subprocess
import webbrowser
import time
import threading
import os
import sys
import socket
import signal
import json
import psutil
import logging
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from datetime import datetime

# Path configurations
WORKSPACE_PATH = Path(__file__).parent.parent
LAUNCHER_PATH = WORKSPACE_PATH / "ETF_Suite_Launcher.py"
PORTFOLIO_BUILDER_PATH = WORKSPACE_PATH / "pages" / "ETF_Portfolio_Builder.py"
ANALYZER_PATH = WORKSPACE_PATH / "pages" / "ETF_Analyzer.py"
CONFIG_DIR = WORKSPACE_PATH / "config"
CONFIG_FILE = CONFIG_DIR / "etf_suite_config.json"
LOGS_DIR = WORKSPACE_PATH / "logs"

# Default port settings
DEFAULT_PORTS = {
    "launcher": 8500,
    "portfolio_builder": 8501,
    "analyzer": 8502
}

# Process tracking
active_processes = {}

# Setup logging
def setup_logging():
    """Set up logging configuration."""
    LOGS_DIR.mkdir(exist_ok=True)
    log_file = LOGS_DIR / f"etf_portal_{datetime.now().strftime('%Y%m%d')}.log"
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger("etf_portal")

logger = setup_logging()

def ensure_config_exists():
    """Ensure config directory and file exist."""
    CONFIG_DIR.mkdir(exist_ok=True)
    
    if not CONFIG_FILE.exists():
        with open(CONFIG_FILE, 'w') as f:
            json.dump({
                "ports": DEFAULT_PORTS,
                "streamlit_path": os.path.join(sys.prefix, "bin", "streamlit")
            }, f, indent=2)

def get_config():
    """Get the configuration from the config file."""
    ensure_config_exists()
    with open(CONFIG_FILE, 'r') as f:
        return json.load(f)

def update_config(key, value):
    """Update a specific configuration value."""
    config = get_config()
    
    # Handle nested configuration like ports.launcher
    if '.' in key:
        main_key, sub_key = key.split('.', 1)
        if main_key not in config:
            config[main_key] = {}
        config[main_key][sub_key] = value
    else:
        config[key] = value
        
    with open(CONFIG_FILE, 'w') as f:
        json.dump(config, f, indent=2)

def cleanup_streamlit_processes():
    """Kill any existing Streamlit processes to prevent conflicts."""
    click.echo("Cleaning up existing Streamlit processes...")
    logger.info("Cleaning up existing Streamlit processes")
    try:
        config = get_config()
        ports = config["ports"]
        processed_pids = set()  # Track PIDs we've already handled
        
        # First, find and kill processes using our target ports
        for port in ports.values():
            try:
                # Find process using the port
                cmd = f"lsof -i :{port} | grep LISTEN | awk '{{print $2}}'"
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
                if result.stdout.strip():
                    pids = result.stdout.strip().split('\n')
                    for pid in pids:
                        pid = int(pid)
                        if pid not in processed_pids:
                            try:
                                os.kill(pid, signal.SIGTERM)
                                logger.info(f"Terminated process {pid} using port {port}")
                                processed_pids.add(pid)
                            except ProcessLookupError:
                                pass
            except Exception as e:
                logger.error(f"Error cleaning up port {port}: {e}")
        
        # Then find and kill any remaining Streamlit processes
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if proc.info['pid'] not in processed_pids and 'streamlit' in ' '.join(proc.info['cmdline'] or []).lower():
                    proc.terminate()
                    logger.info(f"Terminated Streamlit process {proc.info['pid']}")
                    processed_pids.add(proc.info['pid'])
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        
        # Give processes time to shut down
        time.sleep(2)
        
        # Force kill any remaining processes
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if proc.info['pid'] not in processed_pids and 'streamlit' in ' '.join(proc.info['cmdline'] or []).lower():
                    proc.kill()
                    logger.info(f"Force killed Streamlit process {proc.info['pid']}")
                    processed_pids.add(proc.info['pid'])
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
                
    except Exception as e:
        logger.error(f"Error during cleanup: {e}")
        click.echo(f"Warning during cleanup: {e}")

def port_is_available(port):
    """Check if a port is available."""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", port))
            return True
    except socket.error:
        return False

def open_browser(url, delay=3):
    """Open browser after a delay to ensure app is running."""
    time.sleep(delay)
    click.echo(f"Opening browser to {url}")
    webbrowser.open(url)

def start_component(component: str, open_browser_tab=True, background=False) -> Optional[subprocess.Popen]:
    """Start a specific component of the ETF Portal."""
    config = get_config()
    streamlit_path = config.get("streamlit_path", os.path.join(sys.prefix, "bin", "streamlit"))
    ports = config["ports"]
    
    # Check if streamlit exists at the specified path
    if not os.path.exists(streamlit_path):
        error_msg = f"ERROR: Streamlit not found at {streamlit_path}"
        logger.error(error_msg)
        click.echo(error_msg)
        click.echo("Please install streamlit or update the streamlit_path in config")
        return None
    
    # Component-specific configurations
    component_configs = {
        "launcher": {
            "path": LAUNCHER_PATH,
            "port": ports["launcher"],
        },
        "portfolio_builder": {
            "path": PORTFOLIO_BUILDER_PATH,
            "port": ports["portfolio_builder"],
        },
        "analyzer": {
            "path": ANALYZER_PATH,
            "port": ports["analyzer"],
        }
    }
    
    if component not in component_configs:
        error_msg = f"Unknown component: {component}"
        logger.error(error_msg)
        click.echo(error_msg)
        return None
    
    component_config = component_configs[component]
    port = component_config["port"]
    
    if not port_is_available(port):
        error_msg = f"ERROR: Port {port} is in use."
        logger.error(error_msg)
        click.echo(error_msg)
        click.echo(f"Please stop existing service on port {port} or configure a different port.")
        return None
    
    log_file = None
    if background:
        # Create log file for background process
        LOGS_DIR.mkdir(exist_ok=True)
        log_file = LOGS_DIR / f"{component}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        logger.info(f"Starting {component} in background mode. Logs will be written to {log_file}")
    
    click.echo(f"Starting {component} on port {port}..." + (" (background mode)" if background else ""))
    
    # Prepare command
    cmd = [
        streamlit_path, "run", 
        str(component_config["path"]),
        "--server.port", str(port),
        "--server.fileWatcherType", "none"  # Disable file watcher to prevent inotify issues
    ]
    
    # Launch the component
    if background:
        # In background mode, redirect output to log file
        with open(log_file, 'w') as log:
            process = subprocess.Popen(
                cmd,
                stdout=log,
                stderr=log,
                # Make the process independent of the parent
                start_new_session=True
            )
    else:
        # In foreground mode, just run normally
        process = subprocess.Popen(cmd)
    
    # Store process for tracking
    active_processes[component] = {
        "process": process,
        "port": port,
        "pid": process.pid,
        "background": background
    }
    
    # Open browser pointing to the component
    if open_browser_tab:
        thread = threading.Thread(
            target=open_browser,
            args=(f"http://localhost:{port}",)
        )
        thread.daemon = True
        thread.start()
    
    # If running in background, we don't need to wait
    if background:
        logger.info(f"Started {component} in background mode (PID: {process.pid})")
        # Give it a moment to start
        time.sleep(1)
        
        # Check if the process is still running
        if process.poll() is not None:
            error_msg = f"Error: {component} failed to start in background mode"
            logger.error(error_msg)
            click.echo(error_msg)
            return None
    
    return process

def get_streamlit_processes() -> List[Dict]:
    """Get a list of running Streamlit processes."""
    processes = []
    seen_ports = set()  # Track ports we've seen to avoid duplicates
    
    for proc in psutil.process_iter(['pid', 'name', 'cmdline', 'create_time']):
        try:
            cmdline = proc.info['cmdline']
            if not cmdline or 'streamlit' not in ' '.join(cmdline).lower():
                continue
                
            port = None
            component = None
            
            # Extract port from command line
            for i, arg in enumerate(cmdline):
                if arg == '--server.port' and i + 1 < len(cmdline):
                    port = cmdline[i + 1]
                    # Skip if we've already seen this port (likely a duplicate)
                    if port in seen_ports:
                        continue
                    seen_ports.add(port)
            
            # Identify which component based on file path
            for i, arg in enumerate(cmdline):
                if arg == 'run' and i + 1 < len(cmdline):
                    path = cmdline[i + 1]
                    if 'ETF_Suite_Launcher.py' in path:
                        component = 'launcher'
                    elif 'ETF_Portfolio_Builder.py' in path:
                        component = 'portfolio_builder'
                    elif 'ETF_Analyzer.py' in path:
                        component = 'analyzer'
            
            # Only add processes that have a valid port
            if port:
                processes.append({
                    'pid': proc.info['pid'],
                    'port': port,
                    'component': component,
                    'cmdline': ' '.join(cmdline),
                    'create_time': proc.info['create_time']
                })
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    
    # Sort by creation time (newest first)
    processes.sort(key=lambda x: x['create_time'], reverse=True)
    return processes

def stop_component(component=None, pid=None):
    """Stop a specific component or Streamlit process."""
    if pid:
        try:
            process = psutil.Process(pid)
            process.terminate()
            try:
                process.wait(timeout=5)
            except psutil.TimeoutExpired:
                process.kill()
            click.echo(f"Stopped process with PID {pid}")
            logger.info(f"Stopped process with PID {pid}")
            return True
        except psutil.NoSuchProcess:
            click.echo(f"No process found with PID {pid}")
            return False
    elif component:
        # Check active tracked processes first
        if component in active_processes:
            process_info = active_processes[component]
            try:
                process = process_info["process"]
                process.terminate()
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    process.kill()
                click.echo(f"Stopped {component}")
                logger.info(f"Stopped {component}")
                del active_processes[component]
                return True
            except Exception:
                pass
        
        # Fall back to finding by component name in running processes
        processes = get_streamlit_processes()
        for proc in processes:
            if proc['component'] == component:
                return stop_component(pid=proc['pid'])
        
        click.echo(f"No running {component} process found")
        return False

@click.group()
def cli():
    """ETF Portal CLI - Manage the ETF Portal application."""
    pass

@cli.command()
@click.option('--component', type=click.Choice(['launcher', 'portfolio_builder', 'analyzer', 'all']), 
              default='launcher', help='Component to start')
@click.option('--no-browser', is_flag=True, help="Don't open browser automatically")
@click.option('--background', is_flag=True, help="Run in background mode (daemon)")
def start(component, no_browser, background):
    """Start ETF Portal components."""
    if component == 'all':
        # Start launcher first, it will manage the other components
        process = start_component('launcher', not no_browser, background)
        if not process:
            return
    else:
        process = start_component(component, not no_browser, background)
        if not process:
            return
    
    click.echo(f"Started {component}" + (" in background mode" if background else ""))
    
    # In background mode, we just return immediately
    if background:
        return
    
    # In foreground mode, wait for the process
    click.echo("Press Ctrl+C to exit")
    
    # Keep running until interrupted
    try:
        if component == 'all' or component == 'launcher':
            process.wait()
        else:
            # For individual components, we'll just exit
            return
    except KeyboardInterrupt:
        click.echo("\nShutting down...")
        if component == 'all':
            stop_component('launcher')
        else:
            stop_component(component)

@cli.command()
@click.option('--component', type=click.Choice(['launcher', 'portfolio_builder', 'analyzer', 'all']), 
              default='all', help='Component to stop')
@click.option('--pid', type=int, help='Process ID to stop')
def stop(component, pid):
    """Stop ETF Portal components."""
    if pid:
        stop_component(pid=pid)
    elif component == 'all':
        cleanup_streamlit_processes()
        click.echo("Stopped all ETF Portal components")
        logger.info("Stopped all ETF Portal components")
    else:
        stop_component(component)

@cli.command()
@click.option('--component', type=click.Choice(['launcher', 'portfolio_builder', 'analyzer', 'all']), 
              default='all', help='Component to restart')
@click.option('--no-browser', is_flag=True, help="Don't open browser automatically")
@click.option('--background', is_flag=True, help="Run in background mode (daemon)")
def restart(component, no_browser, background):
    """Restart ETF Portal components (stop and then start)."""
    # First stop the components
    if component == 'all':
        cleanup_streamlit_processes()
        click.echo("Stopped all ETF Portal components")
        logger.info("Stopped all ETF Portal components")
    else:
        stop_component(component)
    
    # Give processes time to fully shut down
    time.sleep(2)
    
    # Then start them again
    if component == 'all':
        # Start launcher first, it will manage the other components
        process = start_component('launcher', not no_browser, background)
        if not process:
            return
    else:
        process = start_component(component, not no_browser, background)
        if not process:
            return
    
    click.echo(f"Restarted {component}" + (" in background mode" if background else ""))
    
    # In background mode, we just return immediately
    if background:
        return
    
    # In foreground mode, wait for the process
    click.echo("Press Ctrl+C to exit")
    
    # Keep running until interrupted
    try:
        if component == 'all' or component == 'launcher':
            process.wait()
        else:
            # For individual components, we'll just exit
            return
    except KeyboardInterrupt:
        click.echo("\nShutting down...")
        if component == 'all':
            stop_component('launcher')
        else:
            stop_component(component)

@cli.command()
def status():
    """Check status of ETF Portal components."""
    processes = get_streamlit_processes()
    
    if not processes:
        click.echo("No ETF Portal processes are currently running.")
        return
    
    click.echo("Running ETF Portal processes:")
    for i, proc in enumerate(processes):
        component = proc['component'] or 'unknown'
        port = proc['port'] or 'unknown'
        click.echo(f"{i+1}. {component.upper()} (PID: {proc['pid']}, Port: {port})")

@cli.command()
@click.option('--key', required=True, help='Configuration key to update (e.g., ports.launcher)')
@click.option('--value', required=True, help='New value')
def config(key, value):
    """View or update configuration."""
    try:
        # Convert value to integer if possible
        try:
            value = int(value)
        except ValueError:
            pass
            
        update_config(key, value)
        click.echo(f"Updated {key} to {value}")
        logger.info(f"Updated configuration: {key}={value}")
    except Exception as e:
        error_msg = f"Error updating configuration: {e}"
        logger.error(error_msg)
        click.echo(error_msg)

@cli.command()
def logs():
    """Show recent logs from ETF Portal components."""
    LOGS_DIR.mkdir(exist_ok=True)
    
    log_files = sorted(LOGS_DIR.glob("*.log"), key=os.path.getmtime, reverse=True)
    
    if not log_files:
        click.echo("No log files found.")
        return
    
    click.echo("Recent log files:")
    for i, log_file in enumerate(log_files[:5]):  # Show 5 most recent logs
        size = os.path.getsize(log_file) / 1024  # Size in KB
        modified = datetime.fromtimestamp(os.path.getmtime(log_file)).strftime('%Y-%m-%d %H:%M:%S')
        click.echo(f"{i+1}. {log_file.name} ({size:.1f} KB, last modified: {modified})")
    
    # Show most recent log contents
    if log_files:
        most_recent = log_files[0]
        click.echo(f"\nMost recent log ({most_recent.name}):")
        try:
            # Show last 20 lines of the most recent log
            with open(most_recent, 'r') as f:
                lines = f.readlines()
                for line in lines[-20:]:
                    click.echo(line.strip())
        except Exception as e:
            click.echo(f"Error reading log file: {e}")

def main():
    """Main entry point for the CLI."""
    cli()