#!/bin/bash

#==========================================================================
# PyGuardian Universal Installation Script
# Supports: Standalone, Controller, Agent modes
# Author: SmartSolTech Team
# Version: 2.0
#==========================================================================

set -e  # Exit on any error

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

# Global variables
INSTALL_MODE=""
NON_INTERACTIVE=false
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
INSTALL_DIR="/opt/pyguardian"
SERVICE_USER="pyguardian"
CONFIG_DIR="/etc/pyguardian"
LOG_DIR="/var/log/pyguardian"
DATA_DIR="/var/lib/pyguardian"

# Configuration variables
TELEGRAM_BOT_TOKEN=""
ADMIN_ID=""
CONTROLLER_URL=""
AGENT_TOKEN=""
DB_PATH=""

#==========================================================================
# Helper functions
#==========================================================================

print_header() {
    echo -e "${BLUE}"
    echo "=============================================="
    echo "  PyGuardian $1 Installation"
    echo "=============================================="
    echo -e "${NC}"
}

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

print_error() {
    echo -e "${RED}✗ $1${NC}"
}

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

print_info() {
    echo -e "${BLUE}ℹ $1${NC}"
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root or with sudo"
        exit 1
    fi
}

# Parse command line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --mode=*)
                INSTALL_MODE="${1#*=}"
                shift
                ;;
            --non-interactive)
                NON_INTERACTIVE=true
                shift
                ;;
            --controller-url=*)
                CONTROLLER_URL="${1#*=}"
                shift
                ;;
            --agent-token=*)
                AGENT_TOKEN="${1#*=}"
                shift
                ;;
            --telegram-token=*)
                TELEGRAM_BOT_TOKEN="${1#*=}"
                shift
                ;;
            --admin-id=*)
                ADMIN_ID="${1#*=}"
                shift
                ;;
            -h|--help)
                show_usage
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                show_usage
                exit 1
                ;;
        esac
    done
}

show_usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "OPTIONS:"
    echo "  --mode=MODE              Installation mode: standalone, controller, agent"
    echo "  --non-interactive        Non-interactive installation"
    echo "  --controller-url=URL     Controller URL (for agent mode)"
    echo "  --agent-token=TOKEN      Agent authentication token"
    echo "  --telegram-token=TOKEN   Telegram bot token"
    echo "  --admin-id=ID           Telegram admin ID"
    echo "  -h, --help              Show this help"
}

# Interactive mode selection
select_install_mode() {
    if [[ "$NON_INTERACTIVE" == "true" ]]; then
        return
    fi

    print_info "Выберите режим установки:"
    echo ""
    echo "1) Standalone  - Автономный сервер (все в одном)"
    echo "2) Controller  - Центральный контроллер кластера"
    echo "3) Agent      - Агент для подключения к контроллеру"
    echo ""
    
    while true; do
        read -p "Выберите режим (1-3): " choice
        case $choice in
            1)
                INSTALL_MODE="standalone"
                break
                ;;
            2)
                INSTALL_MODE="controller"
                break
                ;;
            3)
                INSTALL_MODE="agent"
                break
                ;;
            *)
                print_error "Неверный выбор. Введите 1, 2 или 3."
                ;;
        esac
    done
}

# Check system requirements
check_requirements() {
    print_info "Проверка системных требований..."
    
    # Check OS
    if [[ ! -f /etc/os-release ]]; then
        print_error "Unsupported operating system"
        exit 1
    fi
    
    . /etc/os-release
    print_success "OS: $NAME $VERSION_ID"
    
    # Check Python version
    if ! command -v python3 &> /dev/null; then
        print_error "Python 3 is required but not installed"
        exit 1
    fi
    
    PYTHON_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
    print_success "Python version: $PYTHON_VERSION"
    
    # Check if Python version is >= 3.10
    if ! python3 -c 'import sys; exit(0 if sys.version_info >= (3, 10) else 1)'; then
        print_error "Python 3.10+ is required, but $PYTHON_VERSION is installed"
        exit 1
    fi
    
    # Check pip
    if ! command -v pip3 &> /dev/null; then
        print_error "pip3 is required but not installed"
        exit 1
    fi
    print_success "pip3 is available"
    
    # Check firewall
    if command -v iptables &> /dev/null; then
        print_success "iptables is available"
    elif command -v nft &> /dev/null; then
        print_success "nftables is available"
    else
        print_warning "Neither iptables nor nftables found - firewall functionality may be limited"
    fi
}

# Install system dependencies
install_dependencies() {
    print_info "Установка системных зависимостей..."
    
    # Detect package manager
    if command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y python3-pip python3-venv python3-dev build-essential \
                          libssl-dev libffi-dev sqlite3 curl wget systemd
        print_success "Dependencies installed (APT)"
    elif command -v yum &> /dev/null; then
        yum install -y python3-pip python3-devel gcc openssl-devel libffi-devel \
                       sqlite curl wget systemd
        print_success "Dependencies installed (YUM)"
    elif command -v dnf &> /dev/null; then
        dnf install -y python3-pip python3-devel gcc openssl-devel libffi-devel \
                       sqlite curl wget systemd
        print_success "Dependencies installed (DNF)"
    else
        print_error "Unsupported package manager"
        exit 1
    fi
}

# Create system user
create_user() {
    print_info "Создание системного пользователя..."
    
    if ! id "$SERVICE_USER" &>/dev/null; then
        useradd --system --create-home --shell /bin/bash "$SERVICE_USER"
        print_success "User $SERVICE_USER created"
    else
        print_info "User $SERVICE_USER already exists"
    fi
}

# Create directories
create_directories() {
    print_info "Создание директорий..."
    
    mkdir -p "$INSTALL_DIR" "$CONFIG_DIR" "$LOG_DIR" "$DATA_DIR"
    chown -R "$SERVICE_USER:$SERVICE_USER" "$INSTALL_DIR" "$LOG_DIR" "$DATA_DIR"
    chmod 755 "$CONFIG_DIR"
    chmod 750 "$LOG_DIR" "$DATA_DIR"
    
    print_success "Directories created"
}

# Copy application files
copy_files() {
    print_info "Копирование файлов приложения..."
    
    # Copy source code
    cp -r "$PROJECT_DIR/src" "$INSTALL_DIR/"
    cp "$PROJECT_DIR/main.py" "$INSTALL_DIR/"
    cp "$PROJECT_DIR/requirements.txt" "$INSTALL_DIR/"
    
    # Copy configuration template
    if [[ "$INSTALL_MODE" == "standalone" || "$INSTALL_MODE" == "controller" ]]; then
        cp "$PROJECT_DIR/config/config.yaml" "$CONFIG_DIR/config.yaml.template"
    fi
    
    # Set permissions
    chown -R "$SERVICE_USER:$SERVICE_USER" "$INSTALL_DIR"
    chmod +x "$INSTALL_DIR/main.py"
    
    print_success "Files copied"
}

# Install Python dependencies
install_python_deps() {
    print_info "Установка Python зависимостей..."
    
    # Create virtual environment
    sudo -u "$SERVICE_USER" python3 -m venv "$INSTALL_DIR/venv"
    
    # Install dependencies
    sudo -u "$SERVICE_USER" "$INSTALL_DIR/venv/bin/pip" install -r "$INSTALL_DIR/requirements.txt"
    
    print_success "Python dependencies installed"
}

# Configure application based on mode
configure_application() {
    print_info "Настройка приложения..."
    
    case "$INSTALL_MODE" in
        "standalone")
            configure_standalone
            ;;
        "controller")
            configure_controller
            ;;
        "agent")
            configure_agent
            ;;
    esac
}

configure_standalone() {
    print_info "Настройка автономного режима..."
    
    # Get configuration from user
    if [[ "$NON_INTERACTIVE" != "true" ]]; then
        get_telegram_config
    fi
    
    # Create configuration file
    create_standalone_config
    print_success "Standalone configuration created"
}

configure_controller() {
    print_info "Настройка контроллера кластера..."
    
    # Get configuration from user
    if [[ "$NON_INTERACTIVE" != "true" ]]; then
        get_telegram_config
        get_controller_config
    fi
    
    # Create configuration file
    create_controller_config
    print_success "Controller configuration created"
}

configure_agent() {
    print_info "Настройка агента..."
    
    # Get configuration from user
    if [[ "$NON_INTERACTIVE" != "true" ]]; then
        get_agent_config
    fi
    
    # Create configuration file
    create_agent_config
    print_success "Agent configuration created"
}

get_telegram_config() {
    if [[ -z "$TELEGRAM_BOT_TOKEN" ]]; then
        echo ""
        print_info "Настройка Telegram бота:"
        echo "1. Создайте бота у @BotFather"
        echo "2. Получите токен бота"
        echo "3. Узнайте ваш chat ID у @userinfobot"
        echo ""
        read -p "Введите токен Telegram бота: " TELEGRAM_BOT_TOKEN
    fi
    
    if [[ -z "$ADMIN_ID" ]]; then
        read -p "Введите ваш Telegram ID (admin): " ADMIN_ID
    fi
}

get_controller_config() {
    echo ""
    print_info "Дополнительные настройки контроллера:"
    read -p "Порт для API контроллера (по умолчанию 8080): " CONTROLLER_PORT
    CONTROLLER_PORT=${CONTROLLER_PORT:-8080}
    
    read -p "Максимальное количество агентов (по умолчанию 50): " MAX_AGENTS
    MAX_AGENTS=${MAX_AGENTS:-50}
}

get_agent_config() {
    if [[ -z "$CONTROLLER_URL" ]]; then
        read -p "URL контроллера (например, https://controller.example.com:8080): " CONTROLLER_URL
    fi
    
    if [[ -z "$AGENT_TOKEN" ]]; then
        read -p "Токен агента (получите у администратора контроллера): " AGENT_TOKEN
    fi
    
    read -p "Имя агента (по умолчанию: $(hostname)): " AGENT_NAME
    AGENT_NAME=${AGENT_NAME:-$(hostname)}
}

create_standalone_config() {
    cat > "$CONFIG_DIR/config.yaml" <<EOF
# PyGuardian Standalone Configuration
# Generated: $(date)

# Telegram Bot Configuration
telegram:
  bot_token: "$TELEGRAM_BOT_TOKEN"
  admin_id: $ADMIN_ID

# Security & Monitoring Settings  
security:
  max_attempts: 5
  time_window: 60
  unban_time: 3600
  
  # STEALTH SECURITY SETTINGS
  authorized_users:
    - "root"
    - "admin"
    - "ubuntu"
    
  honeypot_users:
    - "test"
    - "guest"
    - "user"
    - "admin123"
    - "backup"
    
  stealth_mode_duration: 300
  compromise_indicators:
    - "suspicious_commands"
    - "unusual_login_times"
    - "multiple_failed_then_success"
    - "honeypot_access"

# Log Monitoring
monitoring:
  auth_log_path: "/var/log/auth.log"
  check_interval: 1.0
  
  failed_patterns:
    - "Failed password"
    - "Invalid user"
    - "authentication failure"
    - "Too many authentication failures"
    - "Failed publickey"

# Firewall Configuration
firewall:
  backend: "iptables"
  chain: "INPUT"
  target: "DROP"

# Storage Configuration
storage:
  database_path: "$DATA_DIR/guardian.db"

# Password Management
passwords:
  password_length: 16
  use_special_chars: true
  password_history_size: 5

# Performance Settings
performance:
  cleanup_interval: 3600
  max_records_age: 604800

# Logging Configuration
logging:
  level: "INFO"
  file: "$LOG_DIR/pyguardian.log"
  max_size: 10485760
  backup_count: 5
EOF

    chown "$SERVICE_USER:$SERVICE_USER" "$CONFIG_DIR/config.yaml"
    chmod 640 "$CONFIG_DIR/config.yaml"
}

create_controller_config() {
    CONTROLLER_PORT=${CONTROLLER_PORT:-8080}
    MAX_AGENTS=${MAX_AGENTS:-50}
    
    cat > "$CONFIG_DIR/config.yaml" <<EOF
# PyGuardian Controller Configuration
# Generated: $(date)

# Operating Mode
mode: "controller"

# Telegram Bot Configuration
telegram:
  bot_token: "$TELEGRAM_BOT_TOKEN"
  admin_id: $ADMIN_ID

# Controller Settings
controller:
  port: $CONTROLLER_PORT
  host: "0.0.0.0"
  max_agents: $MAX_AGENTS
  agent_timeout: 300
  heartbeat_interval: 60

# Cluster Configuration
cluster:
  controller_mode: true
  auto_deployment: true
  agent_auto_update: true

# Security Settings (same as standalone)
security:
  max_attempts: 5
  time_window: 60
  unban_time: 3600
  authorized_users:
    - "root"
    - "admin"
    - "ubuntu"
  honeypot_users:
    - "test"
    - "guest"
    - "user"
    - "admin123"
    - "backup"
  stealth_mode_duration: 300

# Storage Configuration
storage:
  database_path: "$DATA_DIR/controller.db"

# Password Management
passwords:
  password_length: 16
  use_special_chars: true
  password_history_size: 5

# Logging Configuration
logging:
  level: "INFO"
  file: "$LOG_DIR/controller.log"
  max_size: 10485760
  backup_count: 5
EOF

    chown "$SERVICE_USER:$SERVICE_USER" "$CONFIG_DIR/config.yaml"
    chmod 640 "$CONFIG_DIR/config.yaml"
}

create_agent_config() {
    cat > "$CONFIG_DIR/config.yaml" <<EOF
# PyGuardian Agent Configuration
# Generated: $(date)

# Operating Mode
mode: "agent"

# Agent Settings
agent:
  name: "$AGENT_NAME"
  controller_url: "$CONTROLLER_URL"
  token: "$AGENT_TOKEN"
  heartbeat_interval: 60
  reconnect_delay: 30

# Log Monitoring
monitoring:
  auth_log_path: "/var/log/auth.log"
  check_interval: 1.0

# Firewall Configuration
firewall:
  backend: "iptables"
  chain: "INPUT"
  target: "DROP"

# Storage Configuration (local cache)
storage:
  database_path: "$DATA_DIR/agent.db"

# Logging Configuration
logging:
  level: "INFO"
  file: "$LOG_DIR/agent.log"
  max_size: 10485760
  backup_count: 5
EOF

    chown "$SERVICE_USER:$SERVICE_USER" "$CONFIG_DIR/config.yaml"
    chmod 640 "$CONFIG_DIR/config.yaml"
}

# Create systemd service
create_service() {
    print_info "Создание systemd сервиса..."
    
    cat > "/etc/systemd/system/pyguardian.service" <<EOF
[Unit]
Description=PyGuardian Security System
After=network.target
Wants=network-online.target

[Service]
Type=exec
User=$SERVICE_USER
Group=$SERVICE_USER
WorkingDirectory=$INSTALL_DIR
Environment=PATH=$INSTALL_DIR/venv/bin
ExecStart=$INSTALL_DIR/venv/bin/python main.py --config=$CONFIG_DIR/config.yaml
ExecReload=/bin/kill -HUP \$MAINPID
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
SyslogIdentifier=pyguardian

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

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable pyguardian
    
    print_success "Systemd service created and enabled"
}

# Start services
start_services() {
    print_info "Запуск сервисов..."
    
    # Start PyGuardian
    systemctl start pyguardian
    
    # Check status
    if systemctl is-active --quiet pyguardian; then
        print_success "PyGuardian service started successfully"
    else
        print_error "Failed to start PyGuardian service"
        print_info "Check logs with: journalctl -u pyguardian -f"
        exit 1
    fi
}

# Show final information
show_completion_info() {
    print_header "Installation Complete"
    
    echo -e "${GREEN}"
    echo "✓ PyGuardian успешно установлен в режиме: $INSTALL_MODE"
    echo ""
    
    case "$INSTALL_MODE" in
        "standalone")
            echo "Конфигурация:"
            echo "  - Telegram бот настроен"
            echo "  - Мониторинг auth.log активен"
            echo "  - Firewall интеграция готова"
            ;;
        "controller")
            echo "Конфигурация контроллера:"
            echo "  - API порт: ${CONTROLLER_PORT:-8080}"
            echo "  - Максимум агентов: ${MAX_AGENTS:-50}"
            echo "  - Telegram управление готово"
            ;;
        "agent")
            echo "Конфигурация агента:"
            echo "  - Контроллер: $CONTROLLER_URL"
            echo "  - Имя агента: $AGENT_NAME"
            echo "  - Подключение к кластеру готово"
            ;;
    esac
    
    echo -e "${NC}"
    echo ""
    print_info "Полезные команды:"
    echo "  systemctl status pyguardian   # Проверить статус"
    echo "  systemctl restart pyguardian  # Перезапустить"
    echo "  journalctl -u pyguardian -f   # Просмотр логов"
    echo ""
    print_info "Файлы конфигурации:"
    echo "  $CONFIG_DIR/config.yaml"
    echo ""
    print_info "Логи:"
    echo "  $LOG_DIR/"
    echo ""
    
    if [[ "$INSTALL_MODE" == "controller" ]]; then
        echo -e "${YELLOW}"
        echo "⚠ Не забудьте:"
        echo "  1. Открыть порт ${CONTROLLER_PORT:-8080} в firewall"
        echo "  2. Настроить SSL сертификат для HTTPS"
        echo "  3. Добавить агенты через Telegram команды"
        echo -e "${NC}"
    fi
    
    if [[ "$INSTALL_MODE" == "agent" ]]; then
        echo -e "${YELLOW}"
        echo "⚠ Примечание:"
        echo "  Агент будет подключаться к контроллеру автоматически"
        echo "  Проверьте статус подключения в логах"
        echo -e "${NC}"
    fi
}

#==========================================================================
# Main installation flow
#==========================================================================

main() {
    # Check if root
    check_root
    
    # Parse command line arguments
    parse_args "$@"
    
    # Show header
    print_header "Universal"
    
    # Select installation mode if not provided
    if [[ -z "$INSTALL_MODE" ]]; then
        select_install_mode
    fi
    
    print_info "Режим установки: $INSTALL_MODE"
    
    # Perform installation steps
    check_requirements
    install_dependencies
    create_user
    create_directories
    copy_files
    install_python_deps
    configure_application
    create_service
    start_services
    
    # Show completion information
    show_completion_info
    
    print_success "Установка завершена успешно!"
}

# Run main function with all arguments
main "$@"