#!/bin/bash

#==========================================================================
# PyGuardian Docker Installation Script
# Supports containerized deployment for Controller and Agent modes
# Author: SmartSolTech Team
# Version: 2.0
#==========================================================================

set -e

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

# Global variables
INSTALL_MODE=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
DOCKER_COMPOSE_VERSION="2.20.0"

# Configuration variables
TELEGRAM_BOT_TOKEN=""
ADMIN_ID=""
CONTROLLER_URL=""
AGENT_TOKEN=""
CONTROLLER_PORT="8080"

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

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

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

print_error() {
    echo -e "${RED}✗ $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
                ;;
            --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
                ;;
            --port=*)
                CONTROLLER_PORT="${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: controller, agent"
    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 "  --port=PORT             Controller port (default: 8080)"
    echo "  -h, --help              Show this help"
}

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

# Check Docker requirements
check_docker_requirements() {
    print_info "Проверка Docker требований..."
    
    # Check if Docker is installed
    if ! command -v docker &> /dev/null; then
        print_info "Docker не установлен. Устанавливаем Docker..."
        install_docker
    else
        print_success "Docker уже установлен: $(docker --version)"
    fi
    
    # Check if Docker Compose is installed
    if ! command -v docker-compose &> /dev/null; then
        print_info "Docker Compose не установлен. Устанавливаем..."
        install_docker_compose
    else
        print_success "Docker Compose уже установлен: $(docker-compose --version)"
    fi
    
    # Start Docker service
    systemctl start docker
    systemctl enable docker
    print_success "Docker service started and enabled"
}

# Install Docker
install_docker() {
    print_info "Установка Docker..."
    
    # Install prerequisites
    if command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y ca-certificates curl gnupg
        
        # Add Docker's official GPG key
        install -m 0755 -d /etc/apt/keyrings
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
        chmod a+r /etc/apt/keyrings/docker.gpg
        
        # Add repository
        echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
        
        apt-get update
        apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
        
    elif command -v yum &> /dev/null; then
        yum install -y yum-utils
        yum-config-manager --add-repo https://download.docker.com/linux/centos/docker-ce.repo
        yum install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
        
    else
        print_error "Unsupported package manager for Docker installation"
        exit 1
    fi
    
    print_success "Docker installed successfully"
}

# Install Docker Compose
install_docker_compose() {
    print_info "Установка Docker Compose..."
    
    curl -L "https://github.com/docker/compose/releases/download/v${DOCKER_COMPOSE_VERSION}/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
    chmod +x /usr/local/bin/docker-compose
    
    print_success "Docker Compose installed successfully"
}

# Get configuration for controller
get_controller_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
    
    read -p "Порт для API контроллера (по умолчанию $CONTROLLER_PORT): " input_port
    CONTROLLER_PORT=${input_port:-$CONTROLLER_PORT}
}

# Get configuration for agent
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 Dockerfile for controller
create_controller_dockerfile() {
    print_info "Создание Dockerfile для контроллера..."
    
    mkdir -p controller
    
    cat > controller/Dockerfile <<EOF
# PyGuardian Controller Docker Image
FROM python:3.11-slim

LABEL maintainer="SmartSolTech Team"
LABEL description="PyGuardian Security Controller"
LABEL version="2.0"

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    iptables \\
    nftables \\
    curl \\
    sqlite3 \\
    && rm -rf /var/lib/apt/lists/*

# Create app user
RUN useradd --create-home --shell /bin/bash pyguardian

# Set working directory
WORKDIR /app

# Copy requirements first for better caching
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY src/ ./src/
COPY main.py .

# Create necessary directories
RUN mkdir -p /var/lib/pyguardian /var/log/pyguardian /etc/pyguardian
RUN chown -R pyguardian:pyguardian /var/lib/pyguardian /var/log/pyguardian /app

# Copy configuration
COPY controller-config.yaml /etc/pyguardian/config.yaml
RUN chown pyguardian:pyguardian /etc/pyguardian/config.yaml

# Switch to app user
USER pyguardian

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \\
  CMD curl -f http://localhost:8080/health || exit 1

# Expose port
EXPOSE 8080

# Start application
CMD ["python", "main.py", "--config", "/etc/pyguardian/config.yaml"]
EOF

    print_success "Controller Dockerfile created"
}

# Create Dockerfile for agent
create_agent_dockerfile() {
    print_info "Создание Dockerfile для агента..."
    
    mkdir -p agent
    
    cat > agent/Dockerfile <<EOF
# PyGuardian Agent Docker Image
FROM python:3.11-slim

LABEL maintainer="SmartSolTech Team"
LABEL description="PyGuardian Security Agent"
LABEL version="2.0"

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    iptables \\
    nftables \\
    curl \\
    sqlite3 \\
    && rm -rf /var/lib/apt/lists/*

# Create app user
RUN useradd --create-home --shell /bin/bash pyguardian

# Set working directory
WORKDIR /app

# Copy requirements first for better caching
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY src/ ./src/
COPY main.py .

# Create necessary directories
RUN mkdir -p /var/lib/pyguardian /var/log/pyguardian /etc/pyguardian
RUN chown -R pyguardian:pyguardian /var/lib/pyguardian /var/log/pyguardian /app

# Copy configuration
COPY agent-config.yaml /etc/pyguardian/config.yaml
RUN chown pyguardian:pyguardian /etc/pyguardian/config.yaml

# Switch to app user
USER pyguardian

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \\
  CMD python -c "import sys; sys.exit(0)"

# Start application
CMD ["python", "main.py", "--config", "/etc/pyguardian/config.yaml"]
EOF

    print_success "Agent Dockerfile created"
}

# Create controller configuration
create_controller_config_file() {
    cat > controller-config.yaml <<EOF
# PyGuardian Controller Docker Configuration
mode: "controller"

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

# Controller Settings
controller:
  port: 8080
  host: "0.0.0.0"
  max_agents: 50
  agent_timeout: 300
  heartbeat_interval: 60

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

# Security Settings
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: "/var/lib/pyguardian/controller.db"

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

# Logging Configuration
logging:
  level: "INFO"
  file: "/var/log/pyguardian/controller.log"
  max_size: 10485760
  backup_count: 5
EOF
}

# Create agent configuration
create_agent_config_file() {
    cat > agent-config.yaml <<EOF
# PyGuardian Agent Docker Configuration
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
storage:
  database_path: "/var/lib/pyguardian/agent.db"

# Logging Configuration
logging:
  level: "INFO"
  file: "/var/log/pyguardian/agent.log"
  max_size: 10485760
  backup_count: 5
EOF
}

# Create docker-compose.yml for controller
create_controller_compose() {
    print_info "Создание docker-compose.yml для контроллера..."
    
    cat > docker-compose.yml <<EOF
version: '3.8'

services:
  pyguardian-controller:
    build:
      context: .
      dockerfile: controller/Dockerfile
    container_name: pyguardian-controller
    restart: unless-stopped
    ports:
      - "${CONTROLLER_PORT}:8080"
    volumes:
      - controller-data:/var/lib/pyguardian
      - controller-logs:/var/log/pyguardian
      - ./controller-config.yaml:/etc/pyguardian/config.yaml:ro
    environment:
      - TELEGRAM_BOT_TOKEN=${TELEGRAM_BOT_TOKEN}
      - ADMIN_ID=${ADMIN_ID}
    networks:
      - pyguardian-net
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 30s

volumes:
  controller-data:
    driver: local
  controller-logs:
    driver: local

networks:
  pyguardian-net:
    driver: bridge
EOF

    print_success "Controller docker-compose.yml created"
}

# Create docker-compose.yml for agent
create_agent_compose() {
    print_info "Создание docker-compose.yml для агента..."
    
    cat > docker-compose.yml <<EOF
version: '3.8'

services:
  pyguardian-agent:
    build:
      context: .
      dockerfile: agent/Dockerfile
    container_name: pyguardian-agent
    restart: unless-stopped
    network_mode: host
    privileged: true
    volumes:
      - /var/log/auth.log:/var/log/auth.log:ro
      - agent-data:/var/lib/pyguardian
      - agent-logs:/var/log/pyguardian
      - ./agent-config.yaml:/etc/pyguardian/config.yaml:ro
    environment:
      - CONTROLLER_URL=${CONTROLLER_URL}
      - AGENT_TOKEN=${AGENT_TOKEN}
      - AGENT_NAME=${AGENT_NAME}
    cap_add:
      - NET_ADMIN
      - SYS_ADMIN
    healthcheck:
      test: ["CMD", "python", "-c", "import sys; sys.exit(0)"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 30s

volumes:
  agent-data:
    driver: local
  agent-logs:
    driver: local
EOF

    print_success "Agent docker-compose.yml created"
}

# Copy application files
copy_app_files() {
    print_info "Копирование файлов приложения..."
    
    # Copy source files
    cp -r "$PROJECT_DIR/src" ./
    cp "$PROJECT_DIR/main.py" ./
    cp "$PROJECT_DIR/requirements.txt" ./
    
    print_success "Application files copied"
}

# Deploy controller
deploy_controller() {
    print_header "Controller"
    
    get_controller_config
    create_controller_config_file
    copy_app_files
    create_controller_dockerfile
    create_controller_compose
    
    print_info "Запуск контроллера..."
    docker-compose up --build -d
    
    print_info "Ожидание запуска контроллера..."
    sleep 10
    
    if docker-compose ps | grep -q "Up"; then
        print_success "Controller successfully deployed and running"
    else
        print_error "Failed to start controller"
        docker-compose logs
        exit 1
    fi
}

# Deploy agent
deploy_agent() {
    print_header "Agent"
    
    get_agent_config
    create_agent_config_file
    copy_app_files
    create_agent_dockerfile
    create_agent_compose
    
    print_info "Запуск агента..."
    docker-compose up --build -d
    
    print_info "Ожидание запуска агента..."
    sleep 10
    
    if docker-compose ps | grep -q "Up"; then
        print_success "Agent successfully deployed and running"
    else
        print_error "Failed to start agent"
        docker-compose logs
        exit 1
    fi
}

# Show completion info
show_docker_completion_info() {
    print_header "Docker Deployment Complete"
    
    echo -e "${GREEN}✓ PyGuardian успешно развернут в Docker (режим: $INSTALL_MODE)${NC}"
    echo ""
    
    print_info "Полезные Docker команды:"
    echo "  docker-compose ps                 # Статус контейнеров"
    echo "  docker-compose logs -f            # Просмотр логов"
    echo "  docker-compose restart            # Перезапуск"
    echo "  docker-compose stop               # Остановка"
    echo "  docker-compose down               # Полная остановка и удаление"
    echo ""
    
    case "$INSTALL_MODE" in
        "controller")
            echo -e "${YELLOW}⚠ Контроллер настроен:${NC}"
            echo "  - API доступен на порту: $CONTROLLER_PORT"
            echo "  - Telegram бот активен"
            echo "  - Веб-интерфейс: http://localhost:$CONTROLLER_PORT"
            echo ""
            echo -e "${YELLOW}Не забудьте:${NC}"
            echo "  1. Открыть порт $CONTROLLER_PORT в firewall"
            echo "  2. Настроить reverse proxy для HTTPS"
            echo "  3. Добавить агенты через Telegram команды"
            ;;
        "agent")
            echo -e "${YELLOW}⚠ Агент настроен:${NC}"
            echo "  - Подключение к: $CONTROLLER_URL"
            echo "  - Имя агента: $AGENT_NAME"
            echo "  - Мониторинг auth.log активен"
            echo ""
            echo -e "${YELLOW}Примечание:${NC}"
            echo "  Агент автоматически подключится к контроллеру"
            echo "  Проверьте статус в логах контейнера"
            ;;
    esac
}

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

main() {
    check_root
    parse_args "$@"
    
    if [[ -z "$INSTALL_MODE" ]]; then
        select_install_mode
    fi
    
    print_info "Режим Docker установки: $INSTALL_MODE"
    
    check_docker_requirements
    
    case "$INSTALL_MODE" in
        "controller")
            deploy_controller
            ;;
        "agent")
            deploy_agent
            ;;
        *)
            print_error "Unsupported mode: $INSTALL_MODE"
            exit 1
            ;;
    esac
    
    show_docker_completion_info
    print_success "Docker развертывание завершено успешно!"
}

main "$@"