#!/bin/bash

# Exit on error
set -e

# Checkpoint file
CHECKPOINT_FILE="/tmp/mailcow_install_progress"
MAILCOW_DIR="/opt/mailcow-dockerized"

# Function to log messages
log_message() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# Function to save checkpoint
save_checkpoint() {
    echo "$1" > "$CHECKPOINT_FILE"
    log_message "Checkpoint saved: $1"
}

# Function to get last checkpoint
get_checkpoint() {
    if [ -f "$CHECKPOINT_FILE" ]; then
        cat "$CHECKPOINT_FILE"
    else
        echo "start"
    fi
}

# Function to check if a step is completed
is_step_completed() {
    local current_checkpoint=$(get_checkpoint)
    local step=$1
    
    case $current_checkpoint in
        "docker_installed"|"fail2ban_configured"|"mailcow_installed"|"complete")
            if [ "$step" = "system_updated" ]; then return 0; fi
            ;;
        "fail2ban_configured"|"mailcow_installed"|"complete")
            if [ "$step" = "docker_installed" ]; then return 0; fi
            ;;
        "mailcow_installed"|"complete")
            if [ "$step" = "fail2ban_configured" ]; then return 0; fi
            ;;
        "complete")
            if [ "$step" = "mailcow_installed" ]; then return 0; fi
            ;;
    esac
    return 1
}

# Check if script is run as root
if [ "$EUID" -ne 0 ]; then 
    echo "Please run as root"
    exit 1
fi

# Start or resume installation
log_message "Starting or resuming installation from checkpoint: $(get_checkpoint)"

# Step 1: Update system
if ! is_step_completed "system_updated"; then
    log_message "Updating system packages..."
    apt-get update && apt-get upgrade -y
    apt-get install -y ca-certificates curl gnupg ufw fail2ban git sudo
    save_checkpoint "system_updated"
fi

# Step 2: Install Docker
if ! is_step_completed "docker_installed"; then
    log_message "Installing Docker..."
    install -m 0755 -d /etc/apt/keyrings
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
    chmod a+r /etc/apt/keyrings/docker.asc

    echo \
      "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] 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

    # Create docker user if not exists
    if ! id "dockeruser" &>/dev/null; then
        useradd -m -s /bin/bash dockeruser
        usermod -aG docker dockeruser
    fi

    save_checkpoint "docker_installed"
fi

# Step 3: Configure fail2ban
if ! is_step_completed "fail2ban_configured"; then
    log_message "Configuring fail2ban..."
    cat > /etc/fail2ban/jail.local << EOL
[DEFAULT]
bantime = 1h
findtime = 10m
maxretry = 3
banaction = ufw
backend = systemd

ignoreip = 127.0.0.1/8 ::1

[sshd]
enabled = true
port = ssh
filter = sshd
logpath = %(sshd_log)s
maxretry = 3
bantime = 1d
findtime = 10m
EOL

    systemctl restart fail2ban
    save_checkpoint "fail2ban_configured"
fi

# Step 4: Install and configure Mailcow
if ! is_step_completed "mailcow_installed"; then
    log_message "Setting up Mailcow..."
    
    # Check if FQDN is already set
    if [ ! -f "/tmp/mailcow_hostname" ]; then
        read -p "Please enter your Fully Qualified Domain Name (FQDN) (e.g., mail.example.com): " MAILCOW_HOSTNAME
        echo "$MAILCOW_HOSTNAME" > /tmp/mailcow_hostname
    else
        MAILCOW_HOSTNAME=$(cat /tmp/mailcow_hostname)
        log_message "Using saved FQDN: $MAILCOW_HOSTNAME"
    fi

    # Create mailcow directory if not exists
    if [ ! -d "$MAILCOW_DIR" ]; then
        mkdir -p "$MAILCOW_DIR"
        chown dockeruser:dockeruser "$MAILCOW_DIR"
        
        # Clone repository
        su - dockeruser -c "cd /opt && git clone https://github.com/mailcow/mailcow-dockerized"
    fi

    # Generate config if not already generated
    if [ ! -f "$MAILCOW_DIR/mailcow.conf" ]; then
        cd "$MAILCOW_DIR"
        echo "$MAILCOW_HOSTNAME" | ./generate_config.sh
        
        echo "Would you like to edit the mailcow configuration? (y/N)"
        read -n1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            nano mailcow.conf
        fi
    fi

    # Start Mailcow services
    log_message "Starting Mailcow services..."
    su - dockeruser -c "cd $MAILCOW_DIR && docker compose pull && docker compose up -d"
    
    save_checkpoint "mailcow_installed"
fi

# Configure UFW if not already configured
if ! ufw status | grep -q "Status: active"; then
    log_message "Configuring firewall..."
    ufw default deny incoming
    ufw default allow outgoing
    ufw allow 22/tcp
    ufw allow 25/tcp
    ufw allow 80/tcp
    ufw allow 110/tcp
    ufw allow 143/tcp
    ufw allow 443/tcp
    ufw allow 465/tcp
    ufw allow 587/tcp
    ufw allow 993/tcp
    ufw allow 995/tcp
    ufw allow 4190/tcp
    ufw --force enable
fi

# Final steps and cleanup
if [ "$(get_checkpoint)" = "mailcow_installed" ]; then
    log_message "Installation completed successfully!"
    rm -f "$CHECKPOINT_FILE" "/tmp/mailcow_hostname"
    save_checkpoint "complete"
    
    # Print access information
    echo "================================================================"
    echo "Mailcow web interface: https://$MAILCOW_HOSTNAME"
    echo "Docker user: dockeruser"
    echo "Mailcow location: $MAILCOW_DIR"
    echo "Firewall status:"
    ufw status numbered
    echo "================================================================"
fi