#!/bin/bash
# =============================================================================
# Audit Hardening Tool - Docker Service Management Script
# =============================================================================
# Script quản lý Docker container cho Audit Hardening Tool
# Hỗ trợ workflow: Build trên máy có internet → Export → Load trên server dịch vụ
# Tương thích: Docker Compose V1 (docker-compose) & V2 (docker compose)
#
# Usage: ./docker-service.sh [command]
#
# Commands:
#   build     - Build Docker image
#   start     - Start container (docker compose up)
#   stop      - Stop container
#   restart   - Restart container
#   status    - Show container status
#   logs      - View container logs
#   export    - Export image to .tar file for offline deployment
#   load      - Load image from .tar file (on server without internet)
#   clean     - Remove container, image, and volumes
#   help      - Show this help message
# =============================================================================

set -e

# Configuration
IMAGE_NAME="audit-hardening-tool"
CONTAINER_NAME="audit-hardening-tool"
EXPORT_DIR="./docker-export"
COMPOSE_FILE="docker-compose.yml"

# Auto-detect Docker Compose version
# V2: "docker compose" (plugin)   →  COMPOSE_CMD="docker compose -f docker-compose.yml"
# V1: "docker-compose" (standalone) →  COMPOSE_CMD="docker-compose -f docker-compose.yml"
detect_compose() {
    if docker compose version >/dev/null 2>&1; then
        COMPOSE_CMD="docker compose -f ${COMPOSE_FILE}"
        COMPOSE_VER="V2 (plugin)"
    elif command -v docker-compose >/dev/null 2>&1; then
        COMPOSE_CMD="docker-compose -f ${COMPOSE_FILE}"
        COMPOSE_VER="V1 (standalone)"
    else
        echo "ERROR: Neither 'docker compose' nor 'docker-compose' found!"
        echo "Please install Docker Compose: https://docs.docker.com/compose/install/"
        exit 1
    fi
}

# Run detection at script start
detect_compose

# Colors
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

# Print header
print_header() {
    echo ""
    echo -e "${CYAN}=========================================${NC}"
    echo -e "${CYAN}  Audit Hardening Tool - Docker Service${NC}"
    echo -e "${CYAN}=========================================${NC}"
    echo -e "${CYAN}  Compose: ${COMPOSE_VER}${NC}"
    echo ""
}

# Print success message
print_ok() {
    echo -e "${GREEN}✓ $1${NC}"
}

# Print error message
print_error() {
    echo -e "${RED}✗ $1${NC}"
}

# Print info message
print_info() {
    echo -e "${BLUE}ℹ $1${NC}"
}

# Print warning message
print_warn() {
    echo -e "${YELLOW}⚠ $1${NC}"
}

# ================================
# BUILD - Build Docker image
# ================================
cmd_build() {
    print_header
    echo -e "${BLUE}Building Docker image...${NC}"
    echo ""

    # Check Dockerfile exists
    if [ ! -f "Dockerfile" ]; then
        print_error "Dockerfile not found in current directory!"
        exit 1
    fi

    $COMPOSE_CMD build --no-cache
    echo ""
    print_ok "Docker image built successfully!"
    
    # Show image info
    echo ""
    echo -e "${CYAN}Image info:${NC}"
    docker images | grep -E "REPOSITORY|${IMAGE_NAME}" || true
}

# ================================
# START - Start container
# ================================
cmd_start() {
    print_header
    echo -e "${BLUE}Starting container...${NC}"
    echo ""

    # Check required files exist
    check_required_files

    # If no Dockerfile present (offline server), skip build
    if [ -f "Dockerfile" ]; then
        $COMPOSE_CMD up -d
    else
        print_info "No Dockerfile found - using pre-loaded image (offline mode)"
        $COMPOSE_CMD up -d --no-build
    fi
    echo ""
    print_ok "Container started successfully!"
    echo ""

    # Show status
    $COMPOSE_CMD ps
    echo ""
    print_info "Access the application at: http://localhost:8888/audit/"
    print_info "View logs with: $0 logs"
}

# ================================
# STOP - Stop container
# ================================
cmd_stop() {
    print_header
    echo -e "${BLUE}Stopping container...${NC}"
    echo ""

    $COMPOSE_CMD down
    echo ""
    print_ok "Container stopped successfully!"
}

# ================================
# RESTART - Restart container
# ================================
cmd_restart() {
    print_header
    echo -e "${BLUE}Restarting container...${NC}"
    echo ""

    $COMPOSE_CMD restart
    echo ""
    print_ok "Container restarted successfully!"
    echo ""
    $COMPOSE_CMD ps
}

# ================================
# STATUS - Show container status
# ================================
cmd_status() {
    print_header
    echo -e "${CYAN}Container status:${NC}"
    echo ""
    $COMPOSE_CMD ps
    echo ""

    # Show resource usage if running
    if docker ps --format "{{.Names}}" | grep -q "$CONTAINER_NAME"; then
        echo -e "${CYAN}Resource usage:${NC}"
        docker stats --no-stream "$CONTAINER_NAME" 2>/dev/null || true
    fi
}

# ================================
# LOGS - View container logs
# ================================
cmd_logs() {
    print_header
    echo -e "${CYAN}Container logs (last 100 lines, follow mode):${NC}"
    echo -e "${YELLOW}Press Ctrl+C to stop following...${NC}"
    echo ""

    $COMPOSE_CMD logs --tail=100 -f
}

# ================================
# EXPORT - Export image to .tar
# ================================
cmd_export() {
    print_header
    echo -e "${BLUE}Exporting Docker image for offline deployment...${NC}"
    echo ""

    # Check if image exists - try to get image name from compose config
    IMAGE_FULL=$($COMPOSE_CMD config --images 2>/dev/null | head -1) || true
    if [ -z "$IMAGE_FULL" ]; then
        # Fallback: try to find image by name pattern
        IMAGE_FULL=$(docker images --format "{{.Repository}}:{{.Tag}}" | grep "${IMAGE_NAME}" | head -1) || true
    fi
    if [ -z "$IMAGE_FULL" ]; then
        IMAGE_FULL="${IMAGE_NAME}:latest"
    fi

    if ! docker image inspect "$IMAGE_FULL" > /dev/null 2>&1; then
        print_warn "Image not found. Building first..."
        cmd_build
        echo ""
        # Re-detect image name after build
        IMAGE_FULL=$(docker images --format "{{.Repository}}:{{.Tag}}" | grep "${IMAGE_NAME}" | head -1) || true
        if [ -z "$IMAGE_FULL" ]; then
            IMAGE_FULL="${IMAGE_NAME}:latest"
        fi
    fi

    # Create export directory
    mkdir -p "$EXPORT_DIR"

    # Generate filename with timestamp
    TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
    EXPORT_FILE="${EXPORT_DIR}/${IMAGE_NAME}_${TIMESTAMP}.tar"

    echo -e "${CYAN}Saving image: ${IMAGE_FULL}${NC}"
    echo -e "${CYAN}Export file: ${EXPORT_FILE}${NC}"
    echo ""

    docker save -o "$EXPORT_FILE" "$IMAGE_FULL"

    # Show file size
    FILE_SIZE=$(du -h "$EXPORT_FILE" | cut -f1)
    echo ""
    print_ok "Image exported successfully!"
    echo ""
    echo -e "${CYAN}Export details:${NC}"
    echo -e "  File: ${EXPORT_FILE}"
    echo -e "  Size: ${FILE_SIZE}"
    echo ""
    echo -e "${YELLOW}=== HƯỚNG DẪN DEPLOY OFFLINE ===${NC}"
    echo ""
    echo "1. Copy các file sau lên server dịch vụ:"
    echo "   - ${EXPORT_FILE}"
    echo "   - docker-compose.yml"
    echo "   - docker-service.sh"
    echo "   - app/             (thư mục source code của ứng dụng)"
    echo "   - config/          (thư mục chứa checklist + config)"
    echo "   - keys/            (thư mục chứa private/public key)"
    echo "   - tools/           (thư mục chứa hardening tools)"
    echo "   - radius_config.json"
    echo "   - email_config.json"
    echo "   - unit_mapping.json"
    echo ""
    echo "2. Trên server dịch vụ, chạy:"
    echo "   ./docker-service.sh load ${EXPORT_FILE}"
    echo "   ./docker-service.sh start"
    echo ""
}

# ================================
# LOAD - Load image from .tar
# ================================
cmd_load() {
    print_header

    TAR_FILE="$1"

    # If no file specified, find the latest .tar in export dir
    if [ -z "$TAR_FILE" ]; then
        if [ -d "$EXPORT_DIR" ]; then
            TAR_FILE=$(ls -t "${EXPORT_DIR}"/*.tar 2>/dev/null | head -1)
        fi
    fi

    if [ -z "$TAR_FILE" ] || [ ! -f "$TAR_FILE" ]; then
        print_error "No .tar file specified or found!"
        echo ""
        echo "Usage: $0 load <path-to-image.tar>"
        echo ""
        echo "Example:"
        echo "  $0 load ./docker-export/audit-hardening-tool_20260408.tar"
        exit 1
    fi

    FILE_SIZE=$(du -h "$TAR_FILE" | cut -f1)
    echo -e "${BLUE}Loading Docker image from: ${TAR_FILE} (${FILE_SIZE})${NC}"
    echo ""

    docker load -i "$TAR_FILE"
    echo ""
    print_ok "Image loaded successfully!"
    echo ""

    # Show loaded image
    echo -e "${CYAN}Available images:${NC}"
    docker images | grep -E "REPOSITORY|${IMAGE_NAME}" || true
    echo ""
    print_info "Now start the container with: $0 start"
}

# ================================
# CLEAN - Remove everything
# ================================
cmd_clean() {
    print_header
    echo -e "${RED}WARNING: This will remove the container, image, and dangling volumes!${NC}"
    echo ""
    read -p "Are you sure? (y/N): " confirm

    if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
        echo "Cancelled."
        exit 0
    fi

    echo ""
    echo -e "${BLUE}Stopping and removing container...${NC}"
    $COMPOSE_CMD down --rmi local --remove-orphans 2>/dev/null || true

    echo ""
    print_ok "Cleanup completed!"
}

# ================================
# CHECK REQUIRED FILES
# ================================
check_required_files() {
    local has_error=false

    if [ ! -d "app" ]; then
        print_error "Directory 'app/' not found! Ensure application source code is available."
        has_error=true
    fi

    if [ ! -d "config" ]; then
        print_error "Directory 'config/' not found!"
        has_error=true
    fi

    if [ ! -d "keys" ]; then
        print_error "Directory 'keys/' not found!"
        has_error=true
    fi

    if [ ! -f "users_config.json" ]; then
        print_warn "users_config.json not found (will use defaults)"
    fi

    if [ ! -f "radius_config.json" ]; then
        print_warn "radius_config.json not found (RADIUS auth disabled)"
    fi

    if [ ! -f "email_config.json" ]; then
        print_warn "email_config.json not found (email notifications disabled)"
    fi

    if [ ! -f "unit_mapping.json" ]; then
        print_warn "unit_mapping.json not found (using default unit mapping)"
    fi

    if [ ! -d "tools" ]; then
        print_warn "Directory 'tools/' not found (hardening tools won't be available)"
    fi

    if $has_error; then
        print_error "Required files missing. Cannot start!"
        exit 1
    fi
}

# ================================
# HELP - Show help message
# ================================
cmd_help() {
    print_header
    echo "Usage: $0 [command] [options]"
    echo ""
    echo -e "${CYAN}Commands:${NC}"
    echo -e "  ${GREEN}build${NC}            Build Docker image"
    echo -e "  ${GREEN}start${NC}            Start container (docker compose up -d)"
    echo -e "  ${GREEN}stop${NC}             Stop container"
    echo -e "  ${GREEN}restart${NC}          Restart container"
    echo -e "  ${GREEN}status${NC}           Show container status & resource usage"
    echo -e "  ${GREEN}logs${NC}             View container logs (follow mode)"
    echo -e "  ${GREEN}export${NC}           Export image to .tar for offline deployment"
    echo -e "  ${GREEN}load${NC} [file.tar]  Load image from .tar file"
    echo -e "  ${GREEN}clean${NC}            Remove container, image, and volumes"
    echo -e "  ${GREEN}help${NC}             Show this help message"
    echo ""
    echo -e "${CYAN}Offline Deployment Workflow:${NC}"
    echo -e "  ${YELLOW}Máy có internet:${NC}"
    echo "    1. $0 build          # Build image"
    echo "    2. $0 export         # Export thành file .tar"
    echo ""
    echo -e "  ${YELLOW}Server dịch vụ (không internet):${NC}"
    echo "    1. Copy file .tar + các file config lên server"
    echo "    2. $0 load file.tar  # Load image từ .tar"
    echo "    3. $0 start          # Chạy container"
    echo ""
    echo -e "${CYAN}Required files for deployment:${NC}"
    echo "  app/                 Application source code"
    echo "  config/              Checklist & config files (ro)"
    echo "  keys/                RSA keys (ro)"
    echo "  tools/               Hardening tools (rw)"
    echo "  users_config.json    User whitelist (ro)"
    echo "  radius_config.json   RADIUS config (ro)"
    echo "  email_config.json    Email config (ro)"
    echo "  unit_mapping.json    Unit mapping config (ro)"
    echo "  docker-compose.yml   Docker Compose config"
    echo ""
}

# ================================
# MAIN - Entry point
# ================================
case "${1:-help}" in
    build)
        cmd_build
        ;;
    start)
        cmd_start
        ;;
    stop)
        cmd_stop
        ;;
    restart)
        cmd_restart
        ;;
    status)
        cmd_status
        ;;
    logs)
        cmd_logs
        ;;
    export)
        cmd_export
        ;;
    load)
        cmd_load "$2"
        ;;
    clean)
        cmd_clean
        ;;
    help|--help|-h)
        cmd_help
        ;;
    *)
        print_error "Unknown command: $1"
        echo ""
        cmd_help
        exit 1
        ;;
esac