#!/bin/bash

# OpenZiti Router Enrollment Script
# Compatible with Ubuntu 22.04, 24.04 and other Linux distributions
# This script automates the router enrollment process using hash key

set -euo pipefail

# Script configuration
SCRIPT_VERSION="1.0.0"
SCRIPT_NAME="OpenZiti Router Enrollment Script"
LOG_FILE="/var/log/ziti-router-enrollment.log"
CONFIG_DIR="/etc/zitirouter"
CERTS_DIR="${CONFIG_DIR}/certs"
ROUTER_CONFIG="${CONFIG_DIR}/router.yaml"
JWT_FILE="${CONFIG_DIR}/enrollment.jwt"
SYSTEMD_SERVICE_FILE="/etc/systemd/system/ziti-router.service"

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

# Default API endpoint (can be overridden)
DEFAULT_API_ENDPOINT="https://backend.zitinexus.com"


# Initialize variables to prevent unbound variable errors
CALLBACK_URL=""
JWT=""
ROUTER_YAML=""
ROUTER_NAME=""
ROUTER_ID=""
TENANT_ID=""
CONTROLLER_ENDPOINT=""
ROLE_ATTRIBUTES=""
HASH_KEY=""
API_ENDPOINT=""
# Example modification (Automated installation , uncomment below)
#API_ENDPOINT="https://backend.zitinexus.com"
#HASH_KEY="your-hash-key-here"
# Logging function
log() {
    local level=$1
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $message" | tee -a "$LOG_FILE" >/dev/null 2>&1 || true
    
    case $level in
        "ERROR")
            echo -e "${RED}[ERROR]${NC} $message" >&2
            ;;
        "SUCCESS")
            echo -e "${GREEN}[SUCCESS]${NC} $message"
            ;;
        "WARNING")
            echo -e "${YELLOW}[WARNING]${NC} $message"
            ;;
        "INFO")
            echo -e "${BLUE}[INFO]${NC} $message"
            ;;
        *)
            echo "$message"
            ;;
    esac
}

# Error handling
error_exit() {
    log "ERROR" "$1"
    exit 1
}

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

# Check system requirements
check_requirements() {
    log "INFO" "Checking system requirements..."
    
    # Check if curl is installed
    if ! command -v curl &> /dev/null; then
        log "INFO" "Installing curl..."
        apt-get update && apt-get install -y curl || error_exit "Failed to install curl"
    fi
    
    # Check if jq is installed
    if ! command -v jq &> /dev/null; then
        log "INFO" "Installing jq..."
        apt-get update && apt-get install -y jq || error_exit "Failed to install jq"
    fi
    
    # Check if systemctl is available
    if ! command -v systemctl &> /dev/null; then
        error_exit "systemctl is required but not available"
    fi
    
    log "SUCCESS" "System requirements check passed"
}

# Install OpenZiti if not present
install_ziti() {
    if command -v ziti &> /dev/null; then
        local ziti_version=$(ziti version 2>/dev/null | head -n1 || echo "unknown")
        log "INFO" "OpenZiti CLI already installed: $ziti_version"
        return 0
    fi
    
    log "INFO" "Installing OpenZiti CLI using package repository..."
    
    # Method 1: Use OpenZiti package repository (your preferred method)
    log "INFO" "Setting up OpenZiti package repository..."
    
    # Step 1: Add GPG key
    if curl -sSLf https://get.openziti.io/tun/package-repos.gpg | gpg --dearmor --output /usr/share/keyrings/openziti.gpg; then
        log "INFO" "GPG key added successfully"
    else
        error_exit "Failed to add OpenZiti GPG key"
    fi
    
    # Step 2: Set proper permissions
    chmod a+r /usr/share/keyrings/openziti.gpg || error_exit "Failed to set GPG key permissions"
    
    # Step 3: Add repository to sources list
    tee /etc/apt/sources.list.d/openziti-release.list >/dev/null <<EOF
deb [signed-by=/usr/share/keyrings/openziti.gpg] https://packages.openziti.org/zitipax-openziti-deb-stable debian main
EOF
    
    if [[ $? -eq 0 ]]; then
        log "INFO" "Repository added successfully"
    else
        error_exit "Failed to add OpenZiti repository"
    fi
    
    # Step 4: Update package list
    log "INFO" "Updating package list..."
    if apt-get update; then
        log "INFO" "Package list updated successfully"
    else
        error_exit "Failed to update package list"
    fi
    
    # Step 5: Install openziti-router package
    log "INFO" "Installing openziti-router package..."
    if apt-get install -y openziti-router; then
        log "SUCCESS" "OpenZiti router package installed successfully"
    else
        log "WARNING" "Failed to install openziti-router package, trying fallback method..."
        
        # Fallback: Try to install just the ziti CLI
        log "INFO" "Attempting to install ziti CLI only..."
        if apt-get install -y ziti; then
            log "SUCCESS" "OpenZiti CLI installed successfully"
        else
            error_exit "Failed to install OpenZiti CLI from package repository"
        fi
    fi
    
    # Verify installation
    if command -v ziti &> /dev/null; then
        local ziti_version=$(ziti version 2>/dev/null | head -n1 || echo "unknown")
        log "SUCCESS" "OpenZiti CLI installed successfully: $ziti_version"
    else
        error_exit "OpenZiti CLI installation failed - command not found after installation"
    fi
}

# Create necessary directories
create_directories() {
    log "INFO" "Creating necessary directories..."
    
    mkdir -p "$CONFIG_DIR" || error_exit "Failed to create config directory"
    mkdir -p "$CERTS_DIR" || error_exit "Failed to create certs directory"
    mkdir -p "$(dirname "$LOG_FILE")" || error_exit "Failed to create log directory"
    
    # Set proper permissions
    chmod 755 "$CONFIG_DIR"
    chmod 700 "$CERTS_DIR"
    
    log "SUCCESS" "Directories created successfully"
}

# Get user input
get_user_input() {
    echo
    echo "=============================================="
    echo "  $SCRIPT_NAME v$SCRIPT_VERSION"
    echo "=============================================="
    echo
    
    # Get API endpoint
    read -p "Enter ZitiNexus Portal API endpoint [$DEFAULT_API_ENDPOINT]: " api_endpoint
    API_ENDPOINT="${api_endpoint:-$DEFAULT_API_ENDPOINT}"
    
    # Validate API endpoint format
    if [[ ! "$API_ENDPOINT" =~ ^https?:// ]]; then
        error_exit "API endpoint must start with http:// or https://"
    fi
    
    # Get hash key
    echo
    echo "Enter the router enrollment hash key from ZitiNexus Portal:"
    echo "(This is a 32-character hash key provided when you created the router enrollment)"
    read -p "Hash Key: " hash_key
    
    # Validate hash key format
    if [[ ! "$hash_key" =~ ^[a-fA-F0-9]{32}$ ]]; then
        error_exit "Invalid hash key format. Expected 32-character hexadecimal string."
    fi
    
    HASH_KEY="$hash_key"
    
    echo
    log "INFO" "Configuration:"
    log "INFO" "  API Endpoint: $API_ENDPOINT"
    log "INFO" "  Hash Key: ${HASH_KEY:0:8}...${HASH_KEY:24:8}"
    echo
}

# Register router with API
register_router() {
    log "INFO" "Registering router with ZitiNexus Portal..."
    
    local api_url="${API_ENDPOINT}/api/router/register"
    local payload="{\"hashKey\":\"$HASH_KEY\"}"
    local response_file=$(mktemp)
    local http_code
    
    # Debug: Show the URL being called
    log "INFO" "API URL: $api_url"
    
    # Make API call with retry logic
    local max_retries=3
    local retry_count=0
    
    while [ $retry_count -lt $max_retries ]; do
        http_code=$(curl -s -w "%{http_code}" -o "$response_file" \
            -X POST \
            -H "Content-Type: application/json" \
            -H "User-Agent: ZitiRouter-EnrollmentScript/$SCRIPT_VERSION" \
            -d "$payload" \
            --connect-timeout 30 \
            --max-time 60 \
            "$api_url" 2>/dev/null || echo "000")
        
        if [[ "$http_code" == "200" ]]; then
            break
        elif [[ "$http_code" == "429" ]]; then
            retry_count=$((retry_count + 1))
            if [ $retry_count -lt $max_retries ]; then
                local wait_time=$((retry_count * 2))
                log "WARNING" "Rate limited. Waiting ${wait_time}s before retry $retry_count/$max_retries..."
                sleep $wait_time
            fi
        else
            break
        fi
    done
    
    # Check response
    if [[ "$http_code" != "200" ]]; then
        local error_msg="API request failed with HTTP $http_code"
        if [[ -f "$response_file" ]]; then
            local api_error=$(jq -r '.error.message // .message // "Unknown error"' "$response_file" 2>/dev/null || echo "Unknown error")
            error_msg="$error_msg: $api_error"
        fi
        rm -f "$response_file"
        error_exit "$error_msg"
    fi
    
    # Parse response
    if ! jq -e '.success' "$response_file" >/dev/null 2>&1; then
        local api_error=$(jq -r '.error.message // .message // "Registration failed"' "$response_file" 2>/dev/null || echo "Registration failed")
        rm -f "$response_file"
        error_exit "Registration failed: $api_error"
    fi
    
    # Extract data from response - Updated to match expected JSON structure
    JWT=$(jq -r '.data.jwt' "$response_file" 2>/dev/null)
    ROUTER_YAML=$(jq -r '.data.routerConfig.yaml' "$response_file" 2>/dev/null)
    ROUTER_NAME=$(jq -r '.data.routerInfo.name' "$response_file" 2>/dev/null)
    CALLBACK_URL=$(jq -r '.data.callbackUrl // ""' "$response_file" 2>/dev/null)
    
    # Extract additional data for proper router configuration
    ROUTER_ID=$(jq -r '.data.routerInfo.id' "$response_file" 2>/dev/null)
    TENANT_ID=$(jq -r '.data.metadata.tenantId' "$response_file" 2>/dev/null)
    CONTROLLER_ENDPOINT=$(jq -r '.data.metadata.controllerEndpoint' "$response_file" 2>/dev/null)
    
    # Extract role attributes as array
    ROLE_ATTRIBUTES=$(jq -r '.data.routerInfo.roleAttributes[]' "$response_file" 2>/dev/null | tr '\n' ' ')
    
    # Store the full response for debugging
    FULL_RESPONSE=$(cat "$response_file")
    
    rm -f "$response_file"
    
    # Validate extracted data
    if [[ -z "$JWT" || "$JWT" == "null" ]]; then
        error_exit "Failed to extract JWT from API response"
    fi
    
    if [[ -z "$ROUTER_YAML" || "$ROUTER_YAML" == "null" ]]; then
        error_exit "Failed to extract router configuration from API response"
    fi
    
    if [[ -z "$ROUTER_NAME" || "$ROUTER_NAME" == "null" ]]; then
        error_exit "Failed to extract router name from API response"
    fi
    
    if [[ -z "$ROUTER_ID" || "$ROUTER_ID" == "null" ]]; then
        error_exit "Failed to extract router ID from API response"
    fi
    
    if [[ -z "$TENANT_ID" || "$TENANT_ID" == "null" ]]; then
        error_exit "Failed to extract tenant ID from API response"
    fi
    
    # Initialize CALLBACK_URL as empty string if null
    if [[ "$CALLBACK_URL" == "null" ]]; then
        CALLBACK_URL=""
    fi
    
    log "SUCCESS" "Router registered successfully: $ROUTER_NAME (ID: $ROUTER_ID)"
    log "INFO" "  Tenant ID: $TENANT_ID"
    log "INFO" "  Controller: $CONTROLLER_ENDPOINT"
    log "INFO" "  Role Attributes: $ROLE_ATTRIBUTES"
}

# Fix router configuration for proper enrollment
fix_router_configuration() {
    log "INFO" "Constructing router configuration from API response..."
    
    # Create a backup of the original config
    cp "$ROUTER_CONFIG" "${ROUTER_CONFIG}.backup" || error_exit "Failed to backup router configuration"
    
    # Use data from API response instead of hardcoding
    log "INFO" "Using API response data to construct proper router configuration"
    
    # Determine controller endpoint - use from API if available, otherwise default
    local ctrl_endpoint="enroll.zitinexus.com:443"
    if [[ -n "$CONTROLLER_ENDPOINT" && "$CONTROLLER_ENDPOINT" != "null" ]]; then
        ctrl_endpoint="$CONTROLLER_ENDPOINT"
    fi
    
    # Add tls: prefix if not present
    if [[ ! "$ctrl_endpoint" =~ ^tls: ]]; then
        ctrl_endpoint="tls:$ctrl_endpoint"
    fi
    
    # Build role attributes section from API response
    local role_attributes_section=""
    if [[ -n "$ROLE_ATTRIBUTES" && "$ROLE_ATTRIBUTES" != " " ]]; then
        role_attributes_section="roleAttributes:"
        for attr in $ROLE_ATTRIBUTES; do
            role_attributes_section="$role_attributes_section"$'\n'"  - \"$attr\""
        done
    else
        role_attributes_section="# No role attributes specified"
    fi
    
    # Get current timestamp if not provided
    local generated_at=$(date -u +"%Y-%m-%dT%H:%M:%S.%3NZ")
    
    # Generate the complete router configuration using API data
    cat > "$ROUTER_CONFIG" << EOF
v: 3

identity:
  cert: /etc/zitirouter/certs/$ROUTER_NAME.cert
  server_cert: /etc/zitirouter/certs/$ROUTER_NAME.server.chain.cert
  key: /etc/zitirouter/certs/$ROUTER_NAME.key
  ca: /etc/zitirouter/certs/$ROUTER_NAME.cas

ctrl:
  endpoint: $ctrl_endpoint

link:
  dialers:
    - binding: transport

listeners:
# bindings of edge and tunnel requires an "edge" section below
  - binding: edge
    address: tls:0.0.0.0:443
    options:
      advertise: 127.0.0.1:443
      connectTimeoutMs: 5000
      getSessionTimeout: 60
  - binding: tunnel
    options:
      mode: host

edge:
  csr:
    country: SG
    province: SG
    locality: Singapore
    organization: Genworx
    organizationalUnit: ZitiNexus
    sans:
      dns:
        - localhost
        - $ROUTER_NAME
      ip:
        - "127.0.0.1"
        - "::1"

# Tenant-specific role attributes
$role_attributes_section

# Router metadata
metadata:
  tenantId: "$TENANT_ID"
  zitiRouterId: "$ROUTER_ID"
  routerType: "private-edge"
  generatedAt: "$generated_at"
  generatedBy: "ZitiNexus"
EOF
    
    chmod 644 "$ROUTER_CONFIG"
    
    log "SUCCESS" "Router configuration constructed from API response"
    log "INFO" "  Original config backed up to: ${ROUTER_CONFIG}.backup"
    log "INFO" "  Router name: $ROUTER_NAME"
    log "INFO" "  Tenant ID: $TENANT_ID"
    log "INFO" "  Ziti Router ID: $ROUTER_ID"
    log "INFO" "  Controller endpoint: $ctrl_endpoint"
    log "INFO" "  Role attributes: $ROLE_ATTRIBUTES"
}

# Save configuration files
save_configuration() {
    log "INFO" "Saving configuration files..."
    
    # Save JWT
    echo "$JWT" > "$JWT_FILE" || error_exit "Failed to save JWT file"
    chmod 600 "$JWT_FILE"
    
    # Save router configuration
    echo "$ROUTER_YAML" > "$ROUTER_CONFIG" || error_exit "Failed to save router configuration"
    chmod 644 "$ROUTER_CONFIG"
    
    log "SUCCESS" "Configuration files saved"
    log "INFO" "  JWT: $JWT_FILE"
    log "INFO" "  Config: $ROUTER_CONFIG"
    
    # Fix the router configuration for proper enrollment
    fix_router_configuration
}

# Enroll router with OpenZiti
enroll_router() {
    log "INFO" "Enrolling router with OpenZiti controller..."
    
    # Run router enrollment (correct command syntax)
    if ziti router enroll --jwt "$JWT_FILE" "$ROUTER_CONFIG" 2>&1 | tee -a "$LOG_FILE"; then
        log "SUCCESS" "Router enrollment completed successfully"
    else
        error_exit "Router enrollment failed"
    fi
    
    # Verify certificates were created
    if [[ ! -f "${CERTS_DIR}/${ROUTER_NAME}.cert" ]]; then
        error_exit "Router certificate not found after enrollment"
    fi
    
    log "SUCCESS" "Router certificates generated successfully"
}

# Create systemd service
create_systemd_service() {
    log "INFO" "Creating systemd service..."
    
    # Create the final router config file path for systemd
    local final_config="/etc/zitirouter/zitirouter.yaml"
    
    # Copy the router config to the expected location for systemd
    cp "$ROUTER_CONFIG" "$final_config" || error_exit "Failed to copy router config to final location"
    chmod 644 "$final_config"
    
    cat > "$SYSTEMD_SERVICE_FILE" << EOF
[Unit]
Description=OpenZiti Router
After=network.target

[Service]
ExecStart=/usr/bin/ziti router run /etc/zitirouter/zitirouter.yaml
Restart=on-failure
User=root
WorkingDirectory=/etc/zitirouter
LimitNOFILE=65536
StandardOutput=append:/var/log/ziti-router.log
StandardError=append:/var/log/ziti-router.log

[Install]
WantedBy=multi-user.target
EOF
    
    # Reload systemd and enable service
    systemctl daemon-reload || error_exit "Failed to reload systemd"
    systemctl enable ziti-router.service || error_exit "Failed to enable ziti-router service"
    
    log "SUCCESS" "Systemd service created and enabled"
    log "INFO" "  Service file: $SYSTEMD_SERVICE_FILE"
    log "INFO" "  Config file: $final_config"
    log "INFO" "  Log file: /var/log/ziti-router.log"
}

# Start router service
start_router() {
    log "INFO" "Starting OpenZiti router service..."
    
    if systemctl start ziti-router.service; then
        log "SUCCESS" "Router service started successfully"
        
        # Wait a moment and check status
        sleep 3
        if systemctl is-active --quiet ziti-router.service; then
            log "SUCCESS" "Router service is running"
        else
            log "WARNING" "Router service may not be running properly"
            log "INFO" "Check status with: systemctl status ziti-router.service"
        fi
    else
        error_exit "Failed to start router service"
    fi
}

# Report enrollment status back to portal
report_status() {
    log "INFO" "Reporting enrollment status to portal..."
    
    if [[ -z "$CALLBACK_URL" || "$CALLBACK_URL" == "null" ]]; then
        log "WARNING" "No callback URL provided, skipping status report"
        return 0
    fi
    
    # Fix callback URL domain mismatch if needed
    # Replace api.zitinexus.com with backend.zitinexus.com to match script's API endpoint
    local fixed_callback_url="$CALLBACK_URL"
    if [[ "$CALLBACK_URL" == *"api.zitinexus.com"* ]]; then
        fixed_callback_url="${CALLBACK_URL//api.zitinexus.com/backend.zitinexus.com}"
        log "INFO" "Fixed callback URL domain: $fixed_callback_url"
    fi
    
    # Get system information
    local hostname=$(hostname)
    local arch=$(uname -m)
    local os=$(lsb_release -d 2>/dev/null | cut -f2 || echo "Linux")
    local ziti_version=$(ziti version 2>/dev/null | head -n1 || echo "unknown")
    
    local payload=$(cat << EOF
{
    "hashKey": "$HASH_KEY",
    "status": "success",
    "routerInfo": {
        "version": "$ziti_version",
        "hostname": "$hostname",
        "arch": "$arch",
        "os": "$os"
    }
}
EOF
)
    
    local response_file=$(mktemp)
    local http_code
    
    http_code=$(curl -s -w "%{http_code}" -o "$response_file" \
        -X POST \
        -H "Content-Type: application/json" \
        -H "User-Agent: ZitiRouter-EnrollmentScript/$SCRIPT_VERSION" \
        -d "$payload" \
        --connect-timeout 30 \
        --max-time 60 \
        "$fixed_callback_url" 2>/dev/null || echo "000")
    
    if [[ "$http_code" == "200" ]]; then
        log "SUCCESS" "Enrollment status reported successfully"
    else
        log "WARNING" "Failed to report enrollment status (HTTP $http_code)"
        if [[ -f "$response_file" ]]; then
            local error_msg=$(jq -r '.error.message // .message // "Unknown error"' "$response_file" 2>/dev/null || echo "Unknown error")
            log "WARNING" "Error: $error_msg"
        fi
    fi
    
    rm -f "$response_file"
}

# Report failure status
report_failure() {
    local error_message="$1"
    
    if [[ -z "$CALLBACK_URL" || "$CALLBACK_URL" == "null" || -z "$HASH_KEY" ]]; then
        return 0
    fi
    
    # Fix callback URL domain mismatch if needed
    local fixed_callback_url="$CALLBACK_URL"
    if [[ "$CALLBACK_URL" == *"api.zitinexus.com"* ]]; then
        fixed_callback_url="${CALLBACK_URL//api.zitinexus.com/backend.zitinexus.com}"
    fi
    
    local payload=$(cat << EOF
{
    "hashKey": "$HASH_KEY",
    "status": "failed",
    "error": "$error_message"
}
EOF
)
    
    curl -s -X POST \
        -H "Content-Type: application/json" \
        -H "User-Agent: ZitiRouter-EnrollmentScript/$SCRIPT_VERSION" \
        -d "$payload" \
        --connect-timeout 10 \
        --max-time 30 \
        "$fixed_callback_url" >/dev/null 2>&1 || true
}

# Cleanup function
cleanup() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        log "ERROR" "Script failed with exit code $exit_code"
        if [[ -n "${HASH_KEY:-}" ]]; then
            report_failure "Router enrollment script failed"
        fi
    fi
}

# Show final status
show_final_status() {
    echo
    echo "=============================================="
    echo "  ENROLLMENT COMPLETED SUCCESSFULLY"
    echo "=============================================="
    echo
    log "SUCCESS" "OpenZiti Router enrollment completed!"
    echo
    echo "Router Information:"
    echo "  Name: $ROUTER_NAME"
    echo "  Config: $ROUTER_CONFIG"
    echo "  Certificates: $CERTS_DIR"
    echo "  Service: ziti-router.service"
    echo
    echo "Useful Commands:"
    echo "  Check status:    systemctl status ziti-router"
    echo "  View logs:       journalctl -u ziti-router -f"
    echo "  Stop router:     systemctl stop ziti-router"
    echo "  Start router:    systemctl start ziti-router"
    echo "  Restart router:  systemctl restart ziti-router"
    echo
    echo "Log file: $LOG_FILE"
    echo
}

# Main execution
main() {
    # Set up error handling
    trap cleanup EXIT
    
    log "INFO" "Starting $SCRIPT_NAME v$SCRIPT_VERSION"
    
    # Check if running as root
    check_root
    
    # Check system requirements
    check_requirements
    
    # Install OpenZiti if needed
    install_ziti
    
    # Create directories
    create_directories
    
    # Get user input
    get_user_input
    
    # Register router with API
    register_router
    
    # Save configuration files
    save_configuration
    
    # Enroll router
    enroll_router
    
    # Create systemd service
    create_systemd_service
    
    # Start router
    start_router
    
    # Report success status
    report_status
    
    # Show final status
    show_final_status
    
    log "SUCCESS" "Router enrollment process completed successfully"
}

# Run main function
main "$@"