#!/bin/bash
#
# SecZim - Official Installation Script
# Copyright (c) 2024 SecZim. All rights reserved.
#
# Usage: curl -fsSL https://get.seczim.com/install.sh | sudo bash -s -- YOUR-LICENSE-KEY
#

set -e

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

# Configuration
SECZIM_VERSION="3.0.0"
SECZIM_API_URL="https://licenses.seczim.com/api/v1"
DOWNLOAD_URL="https://downloads.seczim.com"
INSTALL_DIR="/opt/seczim"
CONFIG_DIR="/etc/seczim"
LOG_DIR="/var/log/seczim"
DATA_DIR="/var/lib/seczim"

# Print functions
print_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }

print_banner() {
    echo ""
    echo -e "${BLUE}========================================================${NC}"
    echo -e "${BLUE}                                                        ${NC}"
    echo -e "${BLUE}        SecZim Installer v$SECZIM_VERSION                     ${NC}"
    echo -e "${BLUE}   Email Security System for Zimbra/Postfix            ${NC}"
    echo -e "${BLUE}                                                        ${NC}"
    echo -e "${BLUE}========================================================${NC}"
    echo ""
}

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

# Detect system
detect_system() {
    print_info "Detecting system..."

    ARCH=$(uname -m)
    case "$ARCH" in
        x86_64) ARCH="amd64" ;;
        aarch64|arm64) ARCH="arm64" ;;
        *) print_error "Unsupported architecture: $ARCH"; exit 1 ;;
    esac

    if [ -f /etc/os-release ]; then
        . /etc/os-release
        DISTRO=$ID
    else
        print_error "Cannot detect distribution"
        exit 1
    fi

    print_success "Detected: $DISTRO ($ARCH)"
}

# Validate license and get download token
validate_license() {
    local LICENSE_KEY=$1

    if [ -z "$LICENSE_KEY" ]; then
        print_error "License key is required"
        echo "Usage: $0 YOUR-LICENSE-KEY"
        echo ""
        echo "Don't have a license? Get one at https://seczim.com/pricing"
        exit 1
    fi

    print_info "Validating license key..."

    # Call license API to validate and get download token
    RESPONSE=$(curl -s -X POST "$SECZIM_API_URL/license/verify" \
        -H "Content-Type: application/json" \
        -d "{
            \"license_key\": \"$LICENSE_KEY\",
            \"hostname\": \"$(hostname)\",
            \"ip\": \"$(hostname -I 2>/dev/null | awk '{print $1}' || echo 'unknown')\",
            \"version\": \"$SECZIM_VERSION\"
        }" 2>/dev/null || echo '{"valid":false}')

    # Check if valid
    if echo "$RESPONSE" | grep -q '"valid":true'; then
        # Extract download token
        DOWNLOAD_TOKEN=$(echo "$RESPONSE" | grep -o '"download_token":"[^"]*"' | cut -d'"' -f4)
        if [ -z "$DOWNLOAD_TOKEN" ]; then
            print_warning "No download token received, trying with license key..."
            DOWNLOAD_TOKEN="$LICENSE_KEY"
        fi
        print_success "License validated"
    else
        # Demo mode for testing
        print_warning "License validation failed, using demo mode"
        DOWNLOAD_TOKEN="$LICENSE_KEY"
    fi

    echo "$LICENSE_KEY" > /tmp/seczim_license.key
    echo "$DOWNLOAD_TOKEN" > /tmp/seczim_download_token
}

# Prompt for admin password
prompt_admin_password() {
    echo ""
    echo -e "${BLUE}========================================================${NC}"
    echo -e "${BLUE}          Admin Dashboard Configuration                 ${NC}"
    echo -e "${BLUE}========================================================${NC}"
    echo ""
    print_info "Set up the password for the admin web dashboard"
    echo ""

    while true; do
        echo -n "Enter admin password (min 8 characters): "
        read -s ADMIN_PASSWORD < /dev/tty
        echo ""

        if [ ${#ADMIN_PASSWORD} -lt 8 ]; then
            print_error "Password must be at least 8 characters"
            continue
        fi

        echo -n "Confirm admin password: "
        read -s ADMIN_PASSWORD_CONFIRM < /dev/tty
        echo ""

        if [ "$ADMIN_PASSWORD" != "$ADMIN_PASSWORD_CONFIRM" ]; then
            print_error "Passwords do not match. Please try again."
            continue
        fi

        break
    done

    echo "$ADMIN_PASSWORD" > /tmp/seczim_admin_password
    print_success "Admin password configured"
    echo ""
}

# Generate bcrypt hash for password
generate_password_hash() {
    local password="$1"
    local hash=""

    # Try Python with bcrypt first
    if command -v python3 &> /dev/null; then
        hash=$(python3 -c "
import bcrypt
password = '''$password'''.encode('utf-8')
salt = bcrypt.gensalt(rounds=10)
hashed = bcrypt.hashpw(password, salt)
print(hashed.decode('utf-8'))
" 2>/dev/null) || true
    fi

    # If bcrypt not available, try installing it
    if [ -z "$hash" ]; then
        echo -e "${BLUE}[INFO]${NC} Installing bcrypt for password hashing..." >&2
        pip3 install bcrypt -q 2>/dev/null || pip install bcrypt -q 2>/dev/null || \
            python3 -m pip install bcrypt -q 2>/dev/null || true

        if command -v python3 &> /dev/null; then
            hash=$(python3 -c "
import bcrypt
password = '''$password'''.encode('utf-8')
salt = bcrypt.gensalt(rounds=10)
hashed = bcrypt.hashpw(password, salt)
print(hashed.decode('utf-8'))
" 2>/dev/null) || true
        fi
    fi

    # Fallback: use Docker to generate hash
    if [ -z "$hash" ] && command -v docker &> /dev/null; then
        echo -e "${BLUE}[INFO]${NC} Using Docker to generate password hash..." >&2
        hash=$(docker run --rm python:3.11-alpine sh -c "pip install bcrypt -q && python3 -c \"import bcrypt; print(bcrypt.hashpw(b'$password', bcrypt.gensalt(rounds=10)).decode())\"" 2>/dev/null) || true
    fi

    echo "$hash"
}

# Install Docker
install_docker() {
    if command -v docker &> /dev/null; then
        print_success "Docker already installed"
        return
    fi

    print_info "Installing Docker..."

    case "$DISTRO" in
        ubuntu|debian)
            apt-get update -qq
            apt-get install -y ca-certificates curl gnupg python3-pip
            curl -fsSL https://get.docker.com | bash
            ;;
        centos|rhel|rocky|almalinux)
            curl -fsSL https://get.docker.com | bash
            yum install -y python3-pip 2>/dev/null || dnf install -y python3-pip 2>/dev/null || true
            systemctl start docker
            systemctl enable docker
            ;;
        *)
            print_error "Unsupported distribution: $DISTRO"
            exit 1
            ;;
    esac

    # Install docker-compose if not present
    if ! command -v docker-compose &> /dev/null; then
        curl -L "https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
        chmod +x /usr/local/bin/docker-compose
    fi

    print_success "Docker installed"
}

# Create directories
create_directories() {
    print_info "Creating directories..."
    mkdir -p "$INSTALL_DIR/web/frontend/dist" "$CONFIG_DIR" "$LOG_DIR" "$DATA_DIR/mysql" "$DATA_DIR/redis"
    print_success "Directories created"
}

# Download protected files
download_protected_file() {
    local url="$1"
    local output="$2"
    local token=$(cat /tmp/seczim_download_token)

    curl -fsSL "$url" \
        -H "Authorization: Bearer $token" \
        -o "$output" 2>/dev/null

    if [ ! -s "$output" ]; then
        print_error "Failed to download: $url"
        return 1
    fi
}

# Download all SecZim files
download_seczim() {
    print_info "Downloading SecZim components..."

    # Download schema
    print_info "  Downloading database schema..."
    download_protected_file "$DOWNLOAD_URL/database/schema.sql" "$INSTALL_DIR/init.sql"

    if ! grep -q "CREATE TABLE" "$INSTALL_DIR/init.sql" 2>/dev/null; then
        print_error "Invalid schema file"
        exit 1
    fi
    print_success "  Schema downloaded ($(grep -c 'CREATE TABLE' "$INSTALL_DIR/init.sql") tables)"

    # Download daemon
    print_info "  Downloading seczim-daemon..."
    download_protected_file "$DOWNLOAD_URL/binaries/seczim-daemon-v$SECZIM_VERSION-linux-$ARCH" "$INSTALL_DIR/seczim-daemon"
    chmod +x "$INSTALL_DIR/seczim-daemon"
    print_success "  Daemon downloaded"

    # Download API server
    print_info "  Downloading seczim-api..."
    download_protected_file "$DOWNLOAD_URL/binaries/seczim-api-v$SECZIM_VERSION-linux-$ARCH" "$INSTALL_DIR/seczim-api"
    chmod +x "$INSTALL_DIR/seczim-api"
    print_success "  API server downloaded"

    # Download frontend
    print_info "  Downloading web frontend..."
    download_protected_file "$DOWNLOAD_URL/binaries/seczim-frontend-v$SECZIM_VERSION.tar.gz" "/tmp/seczim-frontend.tar.gz"
    tar -xzf /tmp/seczim-frontend.tar.gz -C "$INSTALL_DIR/web/frontend/dist" 2>/dev/null
    # Remove macOS extended attribute files
    find "$INSTALL_DIR/web/frontend/dist" -name '._*' -delete 2>/dev/null || true
    rm -f /tmp/seczim-frontend.tar.gz
    print_success "  Frontend downloaded"

    print_success "All components downloaded"
}

# Update schema with admin password
update_schema_password() {
    print_info "Configuring admin password..."

    ADMIN_PASSWORD=$(cat /tmp/seczim_admin_password)
    ADMIN_HASH=$(generate_password_hash "$ADMIN_PASSWORD")

    if [ -z "$ADMIN_HASH" ]; then
        print_warning "Could not generate bcrypt hash"
        print_warning "You will need to set the admin password manually"
        return
    fi

    # Escape special characters for sed
    ADMIN_HASH_ESCAPED=$(echo "$ADMIN_HASH" | sed 's/[&/\]/\\&/g')

    # Replace placeholder in schema
    sed -i "s|\\\$2a\\\$10\\\$PLACEHOLDER_HASH_REPLACE_DURING_INSTALL|$ADMIN_HASH_ESCAPED|g" "$INSTALL_DIR/init.sql"

    print_success "Admin password configured"
}

# Generate configuration files
generate_config() {
    print_info "Generating configuration..."

    LICENSE_KEY=$(cat /tmp/seczim_license.key)
    JWT_SECRET=$(openssl rand -base64 32)
    MYSQL_ROOT_PASSWORD=$(openssl rand -base64 16)

    # Generate docker-compose.yml
    cat > "$INSTALL_DIR/docker-compose.yml" <<EOF
version: '3.8'

services:
  mysql:
    image: mysql:8.0
    container_name: seczim-mysql
    restart: always
    environment:
      MYSQL_ROOT_PASSWORD: $MYSQL_ROOT_PASSWORD
      MYSQL_DATABASE: seczim
      MYSQL_USER: seczim
      MYSQL_PASSWORD: seczim
    volumes:
      - $DATA_DIR/mysql:/var/lib/mysql
      - $INSTALL_DIR/init.sql:/docker-entrypoint-initdb.d/init.sql:ro
    ports:
      - "127.0.0.1:3306:3306"
    networks:
      - seczim-network
    healthcheck:
      test: ["CMD", "mysqladmin", "ping", "-h", "localhost"]
      interval: 10s
      timeout: 5s
      retries: 5

  redis:
    image: redis:7-alpine
    container_name: seczim-redis
    restart: always
    command: redis-server --appendonly yes
    volumes:
      - $DATA_DIR/redis:/data
    ports:
      - "127.0.0.1:6379:6379"
    networks:
      - seczim-network
    healthcheck:
      test: ["CMD", "redis-cli", "ping"]
      interval: 10s
      timeout: 5s
      retries: 5

networks:
  seczim-network:
    driver: bridge
EOF

    # Generate seczim.yaml
    cat > "$CONFIG_DIR/seczim.yaml" <<EOF
# SecZim Configuration
# Generated on $(date)

server:
  host: 0.0.0.0
  port: 10031
  protocol: tcp
  workers: 10
  max_connections: 1000

database:
  driver: mysql
  dsn: seczim:seczim@tcp(127.0.0.1:3306)/seczim?parseTime=true

cache:
  enabled: true
  backend: redis
  redis_url: redis://127.0.0.1:6379/0
  ttl: 3600

api:
  enabled: true
  host: 0.0.0.0
  port: 8880
  jwt_secret: "$JWT_SECRET"
  cors_origins: "*"
  token_expire_hours: 24

logging:
  level: info
  format: json
  output: stdout

metrics:
  enabled: true
  host: 0.0.0.0
  port: 9090
  path: /metrics
EOF

    # Save license info
    cat > "$CONFIG_DIR/license.json" <<EOF
{
    "license_key": "$LICENSE_KEY",
    "installed_at": "$(date -Iseconds)",
    "version": "$SECZIM_VERSION",
    "license_server": "$SECZIM_API_URL",
    "server_hostname": "$(hostname)",
    "server_ip": "$(hostname -I 2>/dev/null | awk '{print $1}' || echo 'unknown')"
}
EOF

    chmod 600 "$CONFIG_DIR/seczim.yaml" "$CONFIG_DIR/license.json"

    print_success "Configuration generated"
}

# Create systemd services
create_systemd_services() {
    print_info "Creating systemd services..."

    # SecZim Daemon service
    cat > /etc/systemd/system/seczim-daemon.service <<EOF
[Unit]
Description=SecZim Policy Daemon
Documentation=https://docs.seczim.com
After=network.target docker.service
Requires=docker.service

[Service]
Type=simple
User=root
WorkingDirectory=$INSTALL_DIR
ExecStart=$INSTALL_DIR/seczim-daemon --config $CONFIG_DIR/seczim.yaml
Restart=always
RestartSec=10
StandardOutput=append:$LOG_DIR/seczim-daemon.log
StandardError=append:$LOG_DIR/seczim-daemon.log

[Install]
WantedBy=multi-user.target
EOF

    # SecZim API service
    cat > /etc/systemd/system/seczim-api.service <<EOF
[Unit]
Description=SecZim API Server
Documentation=https://docs.seczim.com
After=network.target docker.service seczim-daemon.service
Requires=docker.service

[Service]
Type=simple
User=root
WorkingDirectory=$INSTALL_DIR
ExecStart=$INSTALL_DIR/seczim-api --config $CONFIG_DIR/seczim.yaml
Restart=always
RestartSec=10
StandardOutput=append:$LOG_DIR/seczim-api.log
StandardError=append:$LOG_DIR/seczim-api.log

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable seczim-daemon seczim-api

    print_success "Systemd services created"
}

# Install uninstall script
install_uninstall_script() {
    print_info "Installing uninstall script..."

    cat > /usr/sbin/seczim-uninstall <<'UNINSTALL_EOF'
#!/bin/bash
#
# SecZim Uninstall Script
#

set -e

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

print_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }

if [ "$EUID" -ne 0 ]; then
    print_error "This script must be run as root"
    exit 1
fi

echo ""
echo -e "${RED}========================================================${NC}"
echo -e "${RED}              SecZim Uninstaller                        ${NC}"
echo -e "${RED}========================================================${NC}"
echo ""

while true; do
    read -p "Are you sure you want to uninstall SecZim? [y/N]: " REPLY < /dev/tty
    if [[ -z "$REPLY" ]]; then
        continue
    elif [[ $REPLY =~ ^[Yy]$ ]]; then
        break
    else
        print_info "Uninstall cancelled"
        exit 0
    fi
done

echo ""
echo -e "${YELLOW}========================================================${NC}"
echo -e "${YELLOW}                 Data Preservation                      ${NC}"
echo -e "${YELLOW}========================================================${NC}"
echo ""
echo "SecZim stores data in /var/lib/seczim including:"
echo "  - MySQL database (policies, logs, settings)"
echo "  - Redis cache"
echo ""
echo -e "${YELLOW}Do you want to keep this data for a future reinstall?${NC}"
echo ""
echo "  [K] Keep data   - Data preserved for reinstall"
echo "  [D] Delete data - All data permanently deleted"
echo ""

KEEP_DATA=true
while true; do
    read -p "Your choice [K/d]: " DATA_CHOICE < /dev/tty
    if [[ -z "$DATA_CHOICE" ]]; then
        continue
    elif [[ $DATA_CHOICE =~ ^[Kk]$ ]]; then
        print_info "Data will be preserved"
        break
    elif [[ $DATA_CHOICE =~ ^[Dd]$ ]]; then
        echo ""
        while true; do
            read -p "Are you SURE you want to DELETE all data? This cannot be undone! [y/N]: " CONFIRM < /dev/tty
            if [[ -z "$CONFIRM" ]]; then
                continue
            elif [[ $CONFIRM =~ ^[Yy]$ ]]; then
                KEEP_DATA=false
                print_warning "Data will be DELETED"
                break
            else
                print_info "Data will be preserved"
                break
            fi
        done
        break
    else
        print_error "Invalid choice. Please enter K or D"
    fi
done

echo ""

print_info "Stopping SecZim services..."
systemctl stop seczim-api 2>/dev/null || true
systemctl stop seczim-daemon 2>/dev/null || true
systemctl disable seczim-api 2>/dev/null || true
systemctl disable seczim-daemon 2>/dev/null || true

print_info "Stopping Docker containers..."
if [ -f /opt/seczim/docker-compose.yml ]; then
    cd /opt/seczim
    if [ "$KEEP_DATA" = true ]; then
        docker-compose down 2>/dev/null || docker compose down 2>/dev/null || true
    else
        docker-compose down -v --rmi all 2>/dev/null || docker compose down -v --rmi all 2>/dev/null || true
    fi
else
    docker rm -f seczim-mysql seczim-redis 2>/dev/null || true
fi

print_info "Removing Docker images..."
docker rmi mysql:8.0 redis:alpine 2>/dev/null || true

print_info "Removing systemd services..."
rm -f /etc/systemd/system/seczim.service
rm -f /etc/systemd/system/seczim-api.service
rm -f /etc/systemd/system/seczim-daemon.service
systemctl daemon-reload

if [ -d /opt/zimbra ]; then
    print_info "Restoring Zimbra configuration..."
    ZIMBRA_USER=$(stat -c '%U' /opt/zimbra 2>/dev/null || echo "zimbra")
    if [ -f /etc/seczim/zimbra-smtpd-restrictions.backup ]; then
        BACKUP=$(cat /etc/seczim/zimbra-smtpd-restrictions.backup 2>/dev/null)
        if [ -n "$BACKUP" ]; then
            su - "$ZIMBRA_USER" -c "/opt/zimbra/bin/postconf -e \"smtpd_recipient_restrictions = $BACKUP\"" >/dev/null 2>&1 || true
        fi
    fi
    if [ -f /etc/seczim/zimbra-recipient-restrictions.backup ]; then
        cp /etc/seczim/zimbra-recipient-restrictions.backup /opt/zimbra/conf/zmconfigd/smtpd_recipient_restrictions.cf 2>/dev/null || true
    fi
    su - "$ZIMBRA_USER" -c "/opt/zimbra/bin/zmmtactl reload" >/dev/null 2>&1 || true
    print_success "Zimbra configuration restored"
fi

if [ -f /etc/seczim/postfix-main.cf.backup ]; then
    print_info "Restoring Postfix configuration..."
    cp /etc/seczim/postfix-main.cf.backup /etc/postfix/main.cf
    postfix reload 2>/dev/null || true
    print_success "Postfix configuration restored"
fi

print_info "Removing SecZim program files..."
rm -rf /opt/seczim
rm -rf /etc/seczim
rm -rf /var/log/seczim

# Remove legacy binaries (from older installations)
rm -f /usr/local/bin/seczim
rm -f /usr/local/bin/seczim-api
rm -f /usr/local/bin/seczim-daemon

if [ "$KEEP_DATA" = true ]; then
    print_info "Keeping data directory: /var/lib/seczim"
else
    print_info "Removing data directory..."
    rm -rf /var/lib/seczim
    print_success "All data deleted"
fi

rm -f /usr/sbin/seczim-uninstall

echo ""
echo -e "${GREEN}========================================================${NC}"
echo -e "${GREEN}          SecZim Uninstalled Successfully              ${NC}"
echo -e "${GREEN}========================================================${NC}"
echo ""

if [ "$KEEP_DATA" = true ]; then
    echo -e "${BLUE}Your data has been preserved at /var/lib/seczim${NC}"
    echo "To reinstall: curl -fsSL https://get.seczim.com/install.sh | sudo bash -s -- YOUR-LICENSE-KEY"
fi

echo ""
UNINSTALL_EOF

    chmod +x /usr/sbin/seczim-uninstall
    print_success "Uninstall script installed: /usr/sbin/seczim-uninstall"
}

# Start all services
start_services() {
    print_info "Starting services..."

    # Start Docker containers
    cd "$INSTALL_DIR"
    if command -v docker-compose &> /dev/null; then
        docker-compose up -d
    else
        docker compose up -d
    fi

    print_info "Waiting for database to be ready..."
    sleep 15

    # Start SecZim services
    systemctl start seczim-daemon
    sleep 2
    systemctl start seczim-api
    sleep 2

    # Verify services
    if systemctl is-active --quiet seczim-daemon && systemctl is-active --quiet seczim-api; then
        print_success "All services started"
    else
        print_warning "Some services may not have started correctly"
        print_info "Check logs: journalctl -u seczim-daemon -u seczim-api"
    fi
}

# Configure mail server (Zimbra or Postfix) to use SecZim policy server
configure_mail_server() {
    print_info "Configuring mail server integration..."

    # Check if Zimbra is installed
    if [ -d /opt/zimbra ]; then
        print_info "Detected Zimbra installation"
        configure_zimbra
    # Check if standalone Postfix is installed
    elif [ -f /etc/postfix/main.cf ]; then
        print_info "Detected Postfix installation"
        configure_postfix
    else
        print_warning "No Zimbra or Postfix detected"
        print_info "You will need to manually configure your mail server to use:"
        print_info "  check_policy_service inet:127.0.0.1:10031"
        return
    fi
}

# Configure Zimbra to use SecZim
configure_zimbra() {
    local ZIMBRA_USER=$(stat -c '%U' /opt/zimbra 2>/dev/null || echo "zimbra")
    local ZIMBRA_CONF="/opt/zimbra/conf/zmconfigd/smtpd_recipient_restrictions.cf"

    # Backup original configuration
    if [ -f "$ZIMBRA_CONF" ]; then
        cp "$ZIMBRA_CONF" "$CONFIG_DIR/zimbra-recipient-restrictions.backup"
        print_info "Backed up Zimbra configuration"
    fi

    # Also backup current postconf settings
    su - "$ZIMBRA_USER" -c "postconf smtpd_recipient_restrictions" 2>/dev/null > "$CONFIG_DIR/zimbra-smtpd-restrictions.backup" || true

    # Check if policy service already configured
    if grep -q "check_policy_service inet:127.0.0.1:10031" "$ZIMBRA_CONF" 2>/dev/null; then
        print_info "SecZim policy service already configured in Zimbra"
        su - "$ZIMBRA_USER" -c "zmmtactl reload" >/dev/null 2>&1 || true
        print_success "Zimbra MTA reloaded"
        return
    fi

    # Add policy service before 'permit' in the restrictions file
    if [ -f "$ZIMBRA_CONF" ]; then
        # Add check_policy_service before the final 'permit'
        sed -i 's/^permit$/check_policy_service inet:127.0.0.1:10031\npermit/' "$ZIMBRA_CONF"
        print_success "Added SecZim policy service to Zimbra configuration"

        # Reload MTA
        print_info "Reloading Zimbra MTA..."
        su - "$ZIMBRA_USER" -c "zmmtactl reload" >/dev/null 2>&1 || su - "$ZIMBRA_USER" -c "zmmtactl restart" >/dev/null 2>&1 || true
        print_success "Zimbra MTA reloaded"
    else
        print_warning "Could not find Zimbra recipient restrictions config"
        print_info "You may need to manually add: check_policy_service inet:127.0.0.1:10031"
    fi

    # Verify configuration was applied
    sleep 2
    if su - "$ZIMBRA_USER" -c "postconf smtpd_recipient_restrictions" 2>/dev/null | grep -q "check_policy_service inet:127.0.0.1:10031"; then
        print_success "Zimbra is now configured to use SecZim"
    else
        print_warning "Configuration may not have been applied. Please verify manually."
    fi
}

# Configure standalone Postfix to use SecZim
configure_postfix() {
    local POSTFIX_MAIN="/etc/postfix/main.cf"

    # Backup original configuration
    cp "$POSTFIX_MAIN" "$CONFIG_DIR/postfix-main.cf.backup"
    print_info "Backed up Postfix configuration"

    # Check if policy service already configured
    if grep -q "check_policy_service inet:127.0.0.1:10031" "$POSTFIX_MAIN"; then
        print_info "SecZim policy service already configured in Postfix"
        postfix reload 2>/dev/null || true
        return
    fi

    # Get current smtpd_recipient_restrictions
    local CURRENT_RESTRICTIONS=$(postconf -h smtpd_recipient_restrictions 2>/dev/null)

    if [ -n "$CURRENT_RESTRICTIONS" ]; then
        # Add check_policy_service before permit (if permit exists) or at the end
        if echo "$CURRENT_RESTRICTIONS" | grep -q "permit"; then
            # Insert before 'permit'
            NEW_RESTRICTIONS=$(echo "$CURRENT_RESTRICTIONS" | sed 's/permit/check_policy_service inet:127.0.0.1:10031, permit/')
        else
            # Append at the end
            NEW_RESTRICTIONS="$CURRENT_RESTRICTIONS, check_policy_service inet:127.0.0.1:10031"
        fi
        postconf -e "smtpd_recipient_restrictions = $NEW_RESTRICTIONS"
    else
        # Set default restrictions with policy service
        postconf -e "smtpd_recipient_restrictions = permit_mynetworks, permit_sasl_authenticated, reject_unauth_destination, check_policy_service inet:127.0.0.1:10031"
    fi

    print_success "Added SecZim policy service to Postfix configuration"

    # Reload Postfix
    print_info "Reloading Postfix..."
    postfix reload 2>/dev/null || systemctl reload postfix 2>/dev/null || true
    print_success "Postfix reloaded"

    # Verify configuration was applied
    if postconf smtpd_recipient_restrictions 2>/dev/null | grep -q "check_policy_service inet:127.0.0.1:10031"; then
        print_success "Postfix is now configured to use SecZim"
    else
        print_warning "Configuration may not have been applied. Please verify manually."
    fi
}

# Cleanup
cleanup() {
    rm -f /tmp/seczim_license.key
    rm -f /tmp/seczim_admin_password
    rm -f /tmp/seczim_download_token
}

# Print completion
print_completion() {
    local SERVER_IP=$(hostname -I 2>/dev/null | awk '{print $1}' || echo "YOUR-SERVER-IP")

    echo ""
    echo -e "${GREEN}========================================================${NC}"
    echo -e "${GREEN}                                                        ${NC}"
    echo -e "${GREEN}        SecZim Installation Complete!                  ${NC}"
    echo -e "${GREEN}                                                        ${NC}"
    echo -e "${GREEN}========================================================${NC}"
    echo ""
    echo -e "${BLUE}Admin Dashboard:${NC}"
    echo "  URL:      http://$SERVER_IP:8880"
    echo "  Username: admin"
    echo "  Password: (the password you set during installation)"
    echo ""
    echo -e "${BLUE}Services:${NC}"
    echo "  Policy Daemon: systemctl status seczim-daemon"
    echo "  API Server:    systemctl status seczim-api"
    echo "  MySQL:         docker ps | grep seczim-mysql"
    echo "  Redis:         docker ps | grep seczim-redis"
    echo ""
    echo -e "${BLUE}Logs:${NC}"
    echo "  Daemon: $LOG_DIR/seczim-daemon.log"
    echo "  API:    $LOG_DIR/seczim-api.log"
    echo ""
    echo -e "${BLUE}Configuration:${NC}"
    echo "  Config: $CONFIG_DIR/seczim.yaml"
    echo "  Data:   $DATA_DIR"
    echo ""
    echo -e "${YELLOW}Next steps:${NC}"
    echo "  1. Configure your mail server to use SecZim"
    echo "  2. Access the web dashboard to configure policies"
    echo "  3. See docs at https://docs.seczim.com"
    echo ""
}

# Main
main() {
    print_banner

    LICENSE_KEY="${1:-}"

    check_root
    detect_system
    validate_license "$LICENSE_KEY"
    prompt_admin_password
    install_docker
    create_directories
    download_seczim
    update_schema_password
    generate_config
    create_systemd_services
    install_uninstall_script
    start_services
    configure_mail_server

    cleanup

    print_completion
}

main "$@"