#!/usr/bin/env python
"""
ETF Suite CLI

A command-line interface for managing the ETF Suite 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
LAUNCHER_PATH = WORKSPACE_PATH / "ETF_Suite_Launcher.py"
PORTFOLIO_BUILDER_PATH = WORKSPACE_PATH / "pages" / "ETF_Dividend_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
}

# Full path to streamlit executable - may need to be adjusted
STREAMLIT_PATH = "/home/pascal/.local/bin/streamlit"

# 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_suite_{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_suite")

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": STREAMLIT_PATH
            }, 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"]
        
        # Find processes using our target ports
        for port in ports.values():
            cmd = f"lsof -i :{port} | grep LISTEN | awk '{{print $2}}' | xargs kill -9 2>/dev/null || true"
            subprocess.run(cmd, shell=True)
            
        # Find and kill any lingering Streamlit processes
        cmd = "pkill -f streamlit || true"
        subprocess.run(cmd, shell=True)
        
        # Give processes time to shut down
        time.sleep(1)
    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 Suite."""
    config = get_config()
    streamlit_path = config.get("streamlit_path", STREAMLIT_PATH)
    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.address", "0.0.0.0",
        "--server.headless", "true",
        "--browser.gatherUsageStats", "false"
    ]
    
    # Add component-specific options
    if component == "portfolio_builder":
        cmd.extend(["--server.baseUrlPath", "/portfolio"])
    elif component == "analyzer":
        cmd.extend(["--server.baseUrlPath", "/analyzer"])
    
    try:
        if background:
            with open(log_file, 'w') as f:
                process = subprocess.Popen(
                    cmd,
                    stdout=f,
                    stderr=f,
                    start_new_session=True
                )
        else:
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True
            )
        
        active_processes[component] = process
        
        if open_browser_tab and not background:
            url = f"http://localhost:{port}"
            threading.Thread(target=open_browser, args=(url,)).start()
        
        return process
    except Exception as e:
        error_msg = f"Error starting {component}: {e}"
        logger.error(error_msg)
        click.echo(error_msg)
        return None

def get_streamlit_processes() -> List[Dict]:
    """Get information about running Streamlit processes."""
    processes = []
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            if 'streamlit' in ' '.join(proc.info['cmdline'] or []):
                processes.append({
                    'pid': proc.info['pid'],
                    'name': proc.info['name'],
                    'cmdline': ' '.join(proc.info['cmdline'] or [])
                })
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    return processes

def stop_component(component=None, pid=None):
    """Stop a specific component or process."""
    if pid:
        try:
            os.kill(pid, signal.SIGTERM)
            click.echo(f"Stopped process {pid}")
            return
        except ProcessLookupError:
            click.echo(f"Process {pid} not found")
            return
    
    if component == "all":
        cleanup_streamlit_processes()
        return
    
    if component in active_processes:
        process = active_processes[component]
        try:
            process.terminate()
            process.wait(timeout=5)
            click.echo(f"Stopped {component}")
        except subprocess.TimeoutExpired:
            process.kill()
            click.echo(f"Force killed {component}")
        del active_processes[component]
    else:
        click.echo(f"No active process found for {component}")

@click.group()
def cli():
    """ETF Suite CLI - Manage your ETF Suite applications."""
    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 Suite components."""
    if component == "all":
        for comp in ['launcher', 'portfolio_builder', 'analyzer']:
            start_component(comp, not no_browser, background)
    else:
        start_component(component, not no_browser, background)

@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 Suite components."""
    stop_component(component, pid)

@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 Suite components."""
    if component == "all":
        stop_component("all")
        time.sleep(2)
        for comp in ['launcher', 'portfolio_builder', 'analyzer']:
            start_component(comp, not no_browser, background)
    else:
        stop_component(component)
        time.sleep(2)
        start_component(component, not no_browser, background)

@cli.command()
def status():
    """Show status of ETF Suite components."""
    processes = get_streamlit_processes()
    if not processes:
        click.echo("No ETF Suite components are running")
        return
    
    click.echo("Running components:")
    for proc in processes:
        click.echo(f"PID: {proc['pid']}")
        click.echo(f"Command: {proc['cmdline']}")
        click.echo("---")

@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):
    """Update ETF Suite configuration."""
    try:
        # Convert value to appropriate type
        if value.isdigit():
            value = int(value)
        elif value.lower() in ('true', 'false'):
            value = value.lower() == 'true'
        
        update_config(key, value)
        click.echo(f"Updated {key} to {value}")
    except Exception as e:
        click.echo(f"Error updating config: {e}")

@cli.command()
def logs():
    """Show recent logs."""
    try:
        log_files = sorted(LOGS_DIR.glob("*.log"), key=lambda x: x.stat().st_mtime, reverse=True)
        if not log_files:
            click.echo("No log files found")
            return
        
        latest_log = log_files[0]
        click.echo(f"Showing last 20 lines of {latest_log.name}:")
        click.echo("---")
        
        with open(latest_log) as f:
            lines = f.readlines()
            for line in lines[-20:]:
                click.echo(line.strip())
    except Exception as e:
        click.echo(f"Error reading logs: {e}")

if __name__ == '__main__':
    cli()