#!/bin/bash
# GPU Monitoring System Setup Script

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INSTALL_DIR="/home/crhon/gpu_monitoring_system"
SERVICE_NAME="gpu-monitoring"
LOG_FILE="/var/log/gpu_monitoring_setup.log"

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1" | tee -a "$LOG_FILE"
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
    exit 1
}

info() {
    echo -e "${BLUE}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

# Check if running as root for system operations
check_root() {
    if [[ $EUID -eq 0 ]]; then
        return 0
    else
        return 1
    fi
}

# Install system dependencies
install_system_deps() {
    log "Installing system dependencies..."
    
    # Update package list
    sudo apt update
    
    # Install required packages
    sudo apt install -y \
        python3 \
        python3-pip \
        python3-venv \
        python3-pyqt5 \
        python3-pyqt5.qtopengl \
        python3-matplotlib \
        python3-flask \
        python3-flask-cors \
        python3-psutil \
        python3-numpy \
        python3-pandas \
        sqlite3
    
    log "System dependencies installed successfully"
}

# Setup Python virtual environment
setup_venv() {
    log "Setting up Python virtual environment..."
    
    cd "$INSTALL_DIR"
    
    # Create virtual environment
    python3 -m venv venv
    
    # Activate virtual environment
    source venv/bin/activate
    
    # Upgrade pip
    pip install --upgrade pip
    
    # Install Python requirements
    pip install -r requirements.txt
    
    log "Virtual environment setup complete"
}

# Create necessary directories
create_directories() {
    log "Creating necessary directories..."
    
    mkdir -p "$INSTALL_DIR/data"
    mkdir -p "$INSTALL_DIR/logs"
    mkdir -p "$INSTALL_DIR/config"
    mkdir -p "$INSTALL_DIR/backups"
    mkdir -p "$INSTALL_DIR/scripts"
    
    log "Directories created successfully"
}

# Setup systemd service
setup_systemd_service() {
    if ! check_root; then
        warn "Systemd service installation requires root privileges. Skipping..."
        return 1
    fi
    
    log "Setting up systemd service..."
    
    # Copy service file
    sudo cp "$INSTALL_DIR/systemd/$SERVICE_NAME.service" "/etc/systemd/system/$SERVICE_NAME.service"
    
    # Reload systemd
    sudo systemctl daemon-reload
    
    # Enable service
    sudo systemctl enable "$SERVICE_NAME"
    
    log "Systemd service installed and enabled"
}

# Setup autostart for desktop environments
setup_autostart() {
    log "Setting up desktop autostart..."
    
    # Create autostart directory
    mkdir -p "$HOME/.config/autostart"
    
    # Create desktop entry for overlay monitor
    cat > "$HOME/.config/autostart/gpu-monitor-overlay.desktop" << EOF
[Desktop Entry]
Type=Application
Name=GPU Monitor Overlay
Comment=GPU monitoring overlay for desktop
Exec=python3 $INSTALL_DIR/gpu_monitor_desktop.py --display overlay
Icon=video-display
Terminal=false
Hidden=false
NoDisplay=false
X-GNOME-Autostart-enabled=true
EOF
    
    # Create desktop entry for web interface
    cat > "$HOME/.config/autostart/gpu-monitor-web.desktop" << EOF
[Desktop Entry]
Type=Application
Name=GPU Monitor Web
Comment=GPU monitoring web interface
Exec=python3 $INSTALL_DIR/web_interface.py
Icon=video-display
Terminal=false
Hidden=false
NoDisplay=false
X-GNOME-Autostart-enabled=true
EOF
    
    log "Desktop autostart configured"
}

# Setup permissions
setup_permissions() {
    if ! check_root; then
        warn "Permission setup requires root privileges. Skipping..."
        return 1
    fi
    
    log "Setting up permissions..."
    
    # Set proper permissions for GPU access
    sudo usermod -a -G video "$USER"
    sudo usermod -a -G render "$USER"
    
    # Set permissions for log directory
    sudo mkdir -p /var/log
    sudo touch /var/log/gpu_fan_control.log
    sudo chown "$USER":"$USER" /var/log/gpu_fan_control.log
    sudo chmod 644 /var/log/gpu_fan_control.log
    
    log "Permissions configured"
}

# Create convenience scripts
create_scripts() {
    log "Creating convenience scripts..."
    
    # Create start script
    cat > "$INSTALL_DIR/scripts/start.sh" << 'EOF'
#!/bin/bash
cd "$(dirname "$0")/.."
source venv/bin/activate
python3 gpu_fan_controller.py --daemon
EOF
    
    # Create stop script
    cat > "$INSTALL_DIR/scripts/stop.sh" << 'EOF'
#!/bin/bash
pkill -f gpu_fan_controller.py
EOF
    
    # Create status script
    cat > "$INSTALL_DIR/scripts/status.sh" << 'EOF'
#!/bin/bash
echo "=== GPU Monitoring System Status ==="
echo "Service Status:"
if systemctl is-active --quiet gpu-monitoring 2>/dev/null; then
    echo "  ✓ Service is running"
else
    echo "  ✗ Service is not running"
fi

echo ""
echo "Process Status:"
if pgrep -f gpu_fan_controller.py > /dev/null; then
    echo "  ✓ Fan controller process is running"
else
    echo "  ✗ Fan controller process is not running"
fi

echo ""
echo "Web Interface:"
if pgrep -f web_interface.py > /dev/null; then
    echo "  ✓ Web interface is running"
    echo "  URL: http://localhost:5000"
else
    echo "  ✗ Web interface is not running"
fi

echo ""
echo "Log Files:"
echo "  /var/log/gpu_fan_control.log"
echo "  $HOME/gpu_monitoring_system/logs/"
EOF
    
    # Create test script
    cat > "$INSTALL_DIR/scripts/test.sh" << 'EOF'
#!/bin/bash
cd "$(dirname "$0")/.."
source venv/bin/activate

echo "=== Testing GPU Monitoring System ==="
echo ""

echo "1. Testing GPU Detection..."
python3 -c "
from gpu_monitoring import GPUManager
manager = GPUManager()
if manager.initialize():
    print('✓ GPU detection successful')
    gpus = manager.get_gpu_list()
    print(f'  Found {len(gpus)} GPU(s): {gpus}')
else:
    print('✗ GPU detection failed')
"

echo ""
echo "2. Testing Fan Control..."
python3 -c "
from gpu_fan_controller import FanController
controller = FanController()
if controller.initialize():
    print('✓ Fan control initialization successful')
    profiles = controller.get_profiles()
    print(f'  Available profiles: {list(profiles.keys())}')
else:
    print('✗ Fan control initialization failed')
"

echo ""
echo "3. Testing Web Interface..."
python3 -c "
import sys
sys.path.append('.')
try:
    from web_interface import app
    print('✓ Web interface import successful')
except Exception as e:
    print(f'✗ Web interface import failed: {e}')
"

echo ""
echo "=== Test Complete ==="
EOF
    
    # Make scripts executable
    chmod +x "$INSTALL_DIR/scripts/"*.sh
    
    log "Convenience scripts created"
}

# Create backup script
create_backup_script() {
    log "Creating backup script..."
    
    cat > "$INSTALL_DIR/scripts/backup.sh" << 'EOF'
#!/bin/bash
cd "$(dirname "$0")/.."

BACKUP_DIR="backups/$(date +%Y%m%d_%H%M%S)"
mkdir -p "$BACKUP_DIR"

echo "Creating backup in $BACKUP_DIR..."

# Backup configuration
cp -r config "$BACKUP_DIR/"

# Backup database
if [ -f "data/gpu_monitoring.db" ]; then
    cp data/gpu_monitoring.db "$BACKUP_DIR/"
fi

# Backup logs (last 7 days)
find logs/ -name "*.log" -mtime -7 -exec cp {} "$BACKUP_DIR/" \;

# Create backup info
cat > "$BACKUP_DIR/backup_info.txt" << BACKUP_INFO
Backup created: $(date)
System: $(uname -a)
GPU Monitoring System Version: 1.0
BACKUP_INFO

echo "Backup created successfully in $BACKUP_DIR"
EOF
    
    chmod +x "$INSTALL_DIR/scripts/backup.sh"
    
    log "Backup script created"
}

# Setup cron jobs for maintenance
setup_cron_jobs() {
    log "Setting up maintenance cron jobs..."
    
    # Create cron job for database cleanup
    (crontab -l 2>/dev/null; echo "0 2 * * * $INSTALL_DIR/scripts/backup.sh") | crontab -
    
    # Create cron job for log rotation
    (crontab -l 2>/dev/null; echo "0 0 * * 0 find $INSTALL_DIR/logs -name '*.log' -mtime +7 -delete") | crontab -
    
    log "Cron jobs configured"
}

# Main installation function
main() {
    echo "=========================================="
    echo "GPU Monitoring System Setup"
    echo "=========================================="
    echo ""
    
    # Check if running from correct directory
    if [[ ! -f "$SCRIPT_DIR/requirements.txt" ]]; then
        error "Please run this script from the GPU monitoring system directory"
    fi
    
    # Initialize log file
    echo "Setup started at $(date)" > "$LOG_FILE"
    
    log "Starting installation process..."
    
    # Create directories
    create_directories
    
    # Install system dependencies
    install_system_deps
    
    # Setup virtual environment
    setup_venv
    
    # Setup systemd service (if root)
    if check_root; then
        setup_systemd_service
        setup_permissions
    fi
    
    # Setup autostart
    setup_autostart
    
    # Create scripts
    create_scripts
    create_backup_script
    
    # Setup cron jobs
    setup_cron_jobs
    
    log "Installation completed successfully!"
    echo ""
    echo "=========================================="
    echo "Installation Summary"
    echo "=========================================="
    echo "✓ System dependencies installed"
    echo "✓ Python virtual environment created"
    echo "✓ Configuration files created"
    echo "✓ Systemd service configured (if root)"
    echo "✓ Desktop autostart configured"
    echo "✓ Convenience scripts created"
    echo "✓ Backup system configured"
    echo ""
    echo "Next steps:"
    echo "1. Review configuration in config/"
    echo "2. Test the system: ./scripts/test.sh"
    echo "3. Start the service: sudo systemctl start $SERVICE_NAME"
    echo "4. View status: ./scripts/status.sh"
    echo ""
    echo "Log file: $LOG_FILE"
    echo "=========================================="
}

# Handle command line arguments
case "${1:-}" in
    "install")
        main
        ;;
    "uninstall")
        if check_root; then
            log "Removing systemd service..."
            sudo systemctl stop "$SERVICE_NAME" 2>/dev/null || true
            sudo systemctl disable "$SERVICE_NAME" 2>/dev/null || true
            sudo rm -f "/etc/systemd/system/$SERVICE_NAME.service"
            sudo systemctl daemon-reload
        fi
        
        log "Removing autostart entries..."
        rm -f "$HOME/.config/autostart/gpu-monitor-*.desktop"
        
        log "Removing cron jobs..."
        crontab -l 2>/dev/null | grep -v "$INSTALL_DIR" | crontab -
        
        log "Uninstallation completed"
        ;;
    "test")
        cd "$INSTALL_DIR"
        source venv/bin/activate
        ./scripts/test.sh
        ;;
    "status")
        cd "$INSTALL_DIR"
        ./scripts/status.sh
        ;;
    *)
        echo "Usage: $0 {install|uninstall|test|status}"
        echo ""
        echo "Commands:"
        echo "  install   - Install the GPU monitoring system"
        echo "  uninstall - Remove the GPU monitoring system"
        echo "  test      - Test the system components"
        echo "  status    - Show system status"
        exit 1
        ;;
esac