#!/bin/bash

#==============================================================================
# Modbus TCP Proxy - Installation Script
# Description: Automated installation, update, and management script
# Version: 2.0
#==============================================================================

set -e  # Exit on error
set -o pipefail  # Exit on pipe failure

#==============================================================================
# Configuration Variables
#==============================================================================
REPO_URL="https://github.com/Xerolux/Modbus-Tcp-Proxy.git"
BASE_DIR="/opt/Modbus-Tcp-Proxy"
CONFIG_DIR="/etc/Modbus-Tcp-Proxy"
CONFIG_FILE="$CONFIG_DIR/config.yaml"
SERVICE_NAME="modbus_proxy.service"
SERVICE_USER="modbus_proxy"
LOG_DIR="/var/log/modbus-proxy"
INSTALL_LOG="$LOG_DIR/install.log"
BACKUP_DIR="$BASE_DIR/backups"

#==============================================================================
# Color Codes for Output
#==============================================================================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

#==============================================================================
# Logging Functions
#==============================================================================
log() {
    local level="$1"
    shift
    local message="$@"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $message" >> "$INSTALL_LOG" 2>/dev/null || true
}

print_success() {
    echo -e "${GREEN}✓${NC} $1"
    log "INFO" "$1"
}

print_error() {
    echo -e "${RED}✗${NC} $1" >&2
    log "ERROR" "$1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
    log "WARNING" "$1"
}

print_info() {
    echo -e "${CYAN}ℹ${NC} $1"
    log "INFO" "$1"
}

print_header() {
    echo -e "\n${BLUE}===${NC} $1 ${BLUE}===${NC}"
    log "INFO" "=== $1 ==="
}

#==============================================================================
# Help Function
#==============================================================================
show_help() {
    cat << EOF
${BLUE}Modbus TCP Proxy - Installation Script${NC}

${CYAN}Usage:${NC}
    $0 [OPTIONS]

${CYAN}Options:${NC}
    install         Install or update Modbus TCP Proxy (default)
    uninstall       Remove Modbus TCP Proxy completely
    restart         Restart the service
    status          Show service status
    logs            Show service logs
    backup          Create a backup of current installation
    -h, --help      Show this help message

${CYAN}Examples:${NC}
    $0 install      # Install or update
    $0 uninstall    # Remove installation
    $0 status       # Check service status
    $0 logs         # View service logs

${CYAN}Requirements:${NC}
    - Ubuntu/Debian-based Linux distribution
    - Root or sudo privileges
    - Internet connection for package installation

EOF
}

#==============================================================================
# Privilege Check
#==============================================================================
check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root or with sudo privileges"
        echo -e "${YELLOW}Please run: sudo $0 $@${NC}"
        exit 1
    fi
}

#==============================================================================
# Dependency Check
#==============================================================================
check_dependencies() {
    print_header "Checking System Dependencies"

    local missing_deps=()

    # Check for required commands
    for cmd in apt git python3 systemctl; do
        if ! command -v $cmd &> /dev/null; then
            missing_deps+=($cmd)
        fi
    done

    if [ ${#missing_deps[@]} -ne 0 ]; then
        print_error "Missing required dependencies: ${missing_deps[*]}"
        print_info "Installing missing dependencies..."
        apt update
        apt install -y git python3 python3-pip python3-venv systemd
    fi

    print_success "All system dependencies are available"
}

#==============================================================================
# Create System User
#==============================================================================
create_service_user() {
    print_header "Setting Up Service User"

    if id -u "$SERVICE_USER" > /dev/null 2>&1; then
        print_info "Service user '$SERVICE_USER' already exists"
    else
        print_info "Creating system user '$SERVICE_USER'..."
        useradd -r -s /bin/false -d "$BASE_DIR" -c "Modbus Proxy Service" "$SERVICE_USER"
        print_success "Service user created"
    fi
}

#==============================================================================
# Setup Directories
#==============================================================================
setup_directories() {
    print_header "Creating Directory Structure"

    # Create base directories
    for dir in "$BASE_DIR" "$CONFIG_DIR" "$LOG_DIR" "$BACKUP_DIR"; do
        if [ ! -d "$dir" ]; then
            mkdir -p "$dir"
            print_success "Created directory: $dir"
        else
            print_info "Directory exists: $dir"
        fi
    done

    # Set ownership and permissions
    chown -R "$SERVICE_USER":"$SERVICE_USER" "$BASE_DIR" "$LOG_DIR"
    chmod 755 "$BASE_DIR" "$LOG_DIR"
    chmod 750 "$CONFIG_DIR"

    # Ensure install log is writable
    touch "$INSTALL_LOG"
    chown "$SERVICE_USER":"$SERVICE_USER" "$INSTALL_LOG"
    chmod 644 "$INSTALL_LOG"
}

#==============================================================================
# Backup Function
#==============================================================================
create_backup() {
    if [ -d "$BASE_DIR/.git" ]; then
        print_header "Creating Backup"

        local backup_name="backup_$(date +%Y%m%d_%H%M%S)"
        local backup_path="$BACKUP_DIR/$backup_name"

        mkdir -p "$backup_path"

        # Backup application files (excluding venv and .git)
        if [ -f "$BASE_DIR/modbus_tcp_proxy.py" ]; then
            cp -r "$BASE_DIR"/*.py "$backup_path/" 2>/dev/null || true
            cp -r "$BASE_DIR"/requirements.txt "$backup_path/" 2>/dev/null || true
        fi

        # Backup config if exists
        if [ -f "$CONFIG_FILE" ]; then
            cp "$CONFIG_FILE" "$backup_path/config.yaml"
        fi

        print_success "Backup created: $backup_path"

        # Keep only last 5 backups
        ls -dt "$BACKUP_DIR"/backup_* 2>/dev/null | tail -n +6 | xargs rm -rf 2>/dev/null || true
    fi
}

#==============================================================================
# Install/Update Repository
#==============================================================================
install_repository() {
    print_header "Installing/Updating Application"

    if [ -d "$BASE_DIR/.git" ]; then
        print_info "Updating existing repository..."

        # Create backup before update
        create_backup

        cd "$BASE_DIR"
        git fetch origin
        local_commit=$(git rev-parse HEAD)
        remote_commit=$(git rev-parse origin/main 2>/dev/null || git rev-parse origin/master)

        if [ "$local_commit" = "$remote_commit" ]; then
            print_success "Already up to date"
        else
            git pull || {
                print_error "Failed to update repository"
                print_warning "You can manually update with: cd $BASE_DIR && sudo git pull"
                exit 1
            }
            print_success "Repository updated successfully"
        fi
    else
        print_info "Cloning repository..."
        git clone "$REPO_URL" "$BASE_DIR" || {
            print_error "Failed to clone repository"
            exit 1
        }
        print_success "Repository cloned successfully"
    fi

    # Set proper ownership
    chown -R "$SERVICE_USER":"$SERVICE_USER" "$BASE_DIR"
}

#==============================================================================
# Setup Python Virtual Environment
#==============================================================================
setup_python_env() {
    print_header "Setting Up Python Environment"

    # Check Python version
    python_version=$(python3 --version | awk '{print $2}')
    print_info "Python version: $python_version"

    if [ -d "$BASE_DIR/venv" ]; then
        print_info "Updating Python dependencies..."
        "$BASE_DIR/venv/bin/pip" install --upgrade pip setuptools wheel
        "$BASE_DIR/venv/bin/pip" install -r "$BASE_DIR/requirements.txt" --upgrade
    else
        print_info "Creating virtual environment..."
        python3 -m venv "$BASE_DIR/venv"

        print_info "Installing Python dependencies..."
        "$BASE_DIR/venv/bin/pip" install --upgrade pip setuptools wheel
        "$BASE_DIR/venv/bin/pip" install -r "$BASE_DIR/requirements.txt"
    fi

    # Set proper ownership
    chown -R "$SERVICE_USER":"$SERVICE_USER" "$BASE_DIR/venv"

    print_success "Python environment ready"
}

#==============================================================================
# Setup Configuration
#==============================================================================
setup_configuration() {
    print_header "Setting Up Configuration"

    if [ -f "$CONFIG_FILE" ]; then
        print_success "Configuration file already exists: $CONFIG_FILE"

        # Check if there's a new default config to compare
        if [ -f "$BASE_DIR/config.default.yaml" ]; then
            print_info "Comparing with default configuration..."

            # Check if configs differ significantly
            if ! diff -q "$CONFIG_FILE" "$BASE_DIR/config.default.yaml" > /dev/null 2>&1; then
                print_warning "Your config differs from default. Review $BASE_DIR/config.default.yaml for new options."
            fi
        fi
    else
        if [ -f "$BASE_DIR/config.default.yaml" ]; then
            print_info "Creating configuration from template..."
            cp "$BASE_DIR/config.default.yaml" "$CONFIG_FILE"
            print_success "Configuration file created: $CONFIG_FILE"
            print_warning "Please edit $CONFIG_FILE with your settings before starting the service!"
        else
            print_error "No default configuration found!"
            print_info "Please create $CONFIG_FILE manually"
            exit 1
        fi
    fi

    # Set proper permissions
    chown "$SERVICE_USER":"$SERVICE_USER" "$CONFIG_FILE"
    chmod 640 "$CONFIG_FILE"
}

#==============================================================================
# Stop Service
#==============================================================================
stop_service() {
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        print_info "Stopping existing service..."
        systemctl stop "$SERVICE_NAME"
        sleep 2
        print_success "Service stopped"
    fi
}

#==============================================================================
# Create Systemd Service
#==============================================================================
create_systemd_service() {
    print_header "Creating Systemd Service"

    stop_service

    cat > /etc/systemd/system/$SERVICE_NAME << EOF
[Unit]
Description=Modbus TCP Proxy Service
Documentation=https://github.com/Xerolux/Modbus-Tcp-Proxy
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=$SERVICE_USER
Group=$SERVICE_USER
WorkingDirectory=$BASE_DIR
ExecStart=$BASE_DIR/venv/bin/python3 $BASE_DIR/modbus_tcp_proxy.py --config $CONFIG_FILE
ExecReload=/bin/kill -HUP \$MAINPID

# Security settings
NoNewPrivileges=true
PrivateTmp=true
ProtectSystem=strict
ProtectHome=true
ReadWritePaths=$LOG_DIR $CONFIG_DIR

# Restart policy
Restart=always
RestartSec=10
StartLimitInterval=200
StartLimitBurst=5

# Resource limits
LimitNOFILE=65535

# Environment
Environment="PYTHONUNBUFFERED=1"
Environment="PYTHONPATH=$BASE_DIR"

[Install]
WantedBy=multi-user.target
EOF

    print_success "Systemd service file created"
}

#==============================================================================
# Start Service
#==============================================================================
start_service() {
    print_header "Starting Service"

    # Reload systemd
    systemctl daemon-reload

    # Enable service
    systemctl enable "$SERVICE_NAME"
    print_success "Service enabled for auto-start"

    # Start service
    systemctl start "$SERVICE_NAME"

    # Wait a moment for service to start
    sleep 3

    # Check if service is running
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        print_success "Service started successfully"

        # Show service status
        echo ""
        systemctl status "$SERVICE_NAME" --no-pager -l || true

    else
        print_error "Service failed to start"
        print_warning "Check logs with: journalctl -u $SERVICE_NAME -n 50"
        print_warning "Or check: $LOG_DIR/modbus_proxy.log"
        exit 1
    fi
}

#==============================================================================
# Verify Installation
#==============================================================================
verify_installation() {
    print_header "Verifying Installation"

    local errors=0

    # Check if service is running
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        print_success "Service is running"
    else
        print_error "Service is not running"
        ((errors++))
    fi

    # Check if service is enabled
    if systemctl is-enabled --quiet "$SERVICE_NAME"; then
        print_success "Service is enabled for auto-start"
    else
        print_warning "Service is not enabled for auto-start"
    fi

    # Check configuration
    if [ -f "$CONFIG_FILE" ]; then
        print_success "Configuration file exists"
    else
        print_error "Configuration file missing"
        ((errors++))
    fi

    # Check Python environment
    if [ -f "$BASE_DIR/venv/bin/python3" ]; then
        print_success "Python virtual environment exists"
    else
        print_error "Python virtual environment missing"
        ((errors++))
    fi

    # Check main script
    if [ -f "$BASE_DIR/modbus_tcp_proxy.py" ]; then
        print_success "Main application file exists"
    else
        print_error "Main application file missing"
        ((errors++))
    fi

    if [ $errors -eq 0 ]; then
        print_success "All verification checks passed"
        return 0
    else
        print_error "Installation verification found $errors error(s)"
        return 1
    fi
}

#==============================================================================
# Installation Function
#==============================================================================
do_install() {
    print_header "Starting Installation"
    log "INFO" "Installation started by user: $SUDO_USER"

    check_dependencies
    create_service_user
    setup_directories
    install_repository
    setup_python_env
    setup_configuration
    create_systemd_service
    start_service
    verify_installation

    echo ""
    print_header "Installation Complete!"
    echo ""
    print_success "Modbus TCP Proxy has been successfully installed"
    echo ""
    echo -e "${CYAN}Configuration:${NC}    $CONFIG_FILE"
    echo -e "${CYAN}Service Name:${NC}     $SERVICE_NAME"
    echo -e "${CYAN}Log Files:${NC}        $LOG_DIR/"
    echo ""
    echo -e "${YELLOW}Useful Commands:${NC}"
    echo -e "  sudo systemctl status $SERVICE_NAME   # Check service status"
    echo -e "  sudo systemctl restart $SERVICE_NAME  # Restart service"
    echo -e "  sudo journalctl -u $SERVICE_NAME -f   # Follow service logs"
    echo -e "  sudo $0 uninstall                     # Uninstall"
    echo ""
}

#==============================================================================
# Uninstall Function
#==============================================================================
do_uninstall() {
    print_header "Uninstalling Modbus TCP Proxy"

    # Confirm uninstall
    echo -e "${YELLOW}WARNING: This will remove all installed files and the service.${NC}"
    echo -e "${YELLOW}Configuration in $CONFIG_DIR will be preserved.${NC}"
    read -p "Are you sure you want to uninstall? (yes/no): " -r
    echo

    if [[ ! $REPLY =~ ^[Yy][Ee][Ss]$ ]]; then
        print_info "Uninstall cancelled"
        exit 0
    fi

    # Stop and disable service
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        print_info "Stopping service..."
        systemctl stop "$SERVICE_NAME"
    fi

    if systemctl is-enabled --quiet "$SERVICE_NAME" 2>/dev/null; then
        print_info "Disabling service..."
        systemctl disable "$SERVICE_NAME"
    fi

    # Remove service file
    if [ -f "/etc/systemd/system/$SERVICE_NAME" ]; then
        rm -f "/etc/systemd/system/$SERVICE_NAME"
        systemctl daemon-reload
        print_success "Service removed"
    fi

    # Create final backup
    if [ -d "$BASE_DIR" ]; then
        create_backup
    fi

    # Remove application directory
    if [ -d "$BASE_DIR" ]; then
        print_info "Removing application directory..."
        rm -rf "$BASE_DIR"
        print_success "Application directory removed"
    fi

    # Remove log directory
    if [ -d "$LOG_DIR" ]; then
        print_info "Removing log directory..."
        rm -rf "$LOG_DIR"
        print_success "Log directory removed"
    fi

    # Remove service user
    if id -u "$SERVICE_USER" > /dev/null 2>&1; then
        print_info "Removing service user..."
        userdel "$SERVICE_USER" 2>/dev/null || true
        print_success "Service user removed"
    fi

    print_success "Uninstall complete!"
    echo ""
    print_info "Configuration preserved at: $CONFIG_DIR"
    print_info "To remove config: sudo rm -rf $CONFIG_DIR"
    echo ""
}

#==============================================================================
# Service Management Functions
#==============================================================================
do_restart() {
    print_info "Restarting service..."
    systemctl restart "$SERVICE_NAME"
    sleep 2
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        print_success "Service restarted successfully"
        systemctl status "$SERVICE_NAME" --no-pager
    else
        print_error "Service failed to restart"
        exit 1
    fi
}

do_status() {
    systemctl status "$SERVICE_NAME" --no-pager -l
}

do_logs() {
    echo -e "${CYAN}Showing last 50 lines of service logs (Ctrl+C to exit follow mode)${NC}"
    echo ""
    journalctl -u "$SERVICE_NAME" -n 50 -f
}

do_backup() {
    check_root "$@"
    create_backup
    echo ""
    print_info "Available backups:"
    ls -lh "$BACKUP_DIR"/ 2>/dev/null || print_warning "No backups found"
}

#==============================================================================
# Main Function
#==============================================================================
main() {
    # Parse command line arguments
    case "${1:-install}" in
        install)
            check_root "$@"
            do_install
            ;;
        uninstall)
            check_root "$@"
            do_uninstall
            ;;
        restart)
            check_root "$@"
            do_restart
            ;;
        status)
            do_status
            ;;
        logs)
            do_logs
            ;;
        backup)
            do_backup
            ;;
        -h|--help|help)
            show_help
            exit 0
            ;;
        *)
            print_error "Unknown command: $1"
            echo ""
            show_help
            exit 1
            ;;
    esac
}

#==============================================================================
# Script Entry Point
#==============================================================================
main "$@"