#!/usr/bin/env bash

set -euo pipefail

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

# Default values
METRICS_PORT="${1:-9145}"
OPENRESTY_USER="nginx"

# Usage message
usage() {
    echo "Usage: $0 [METRICS_PORT]"
    echo "  METRICS_PORT: Port for Prometheus metrics endpoint (default: 9145)"
    exit 1
}

# Logging functions
log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Cleanup function
cleanup() {
    log_error "Installation failed. Cleaning up..."
    if [[ -d /opt/nginx-lua-prometheus ]]; then
        rm -rf /opt/nginx-lua-prometheus
    fi
    if [[ -f /etc/nginx/conf.d/prometheus-metrics.conf ]]; then
        rm -f /etc/nginx/conf.d/prometheus-metrics.conf
    fi
}

trap cleanup ERR

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

# Validate port argument
if [[ $# -gt 1 ]]; then
    usage
fi

if ! [[ "$METRICS_PORT" =~ ^[0-9]+$ ]] || [[ "$METRICS_PORT" -lt 1024 ]] || [[ "$METRICS_PORT" -gt 65535 ]]; then
    log_error "Invalid port number. Please use a port between 1024-65535"
    exit 1
fi

log_info "[1/8] Detecting distribution and package manager..."

# Auto-detect distribution
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_UPDATE="apt update"
            PKG_INSTALL="apt install -y"
            OPENRESTY_CONF_DIR="/etc/nginx"
            OPENRESTY_USER="www-data"
            ;;
        almalinux|rocky|centos|rhel|ol|fedora)
            PKG_MGR="dnf"
            PKG_UPDATE="dnf update -y"
            PKG_INSTALL="dnf install -y"
            OPENRESTY_CONF_DIR="/etc/nginx"
            OPENRESTY_USER="nginx"
            ;;
        amzn)
            PKG_MGR="yum"
            PKG_UPDATE="yum update -y"
            PKG_INSTALL="yum install -y"
            OPENRESTY_CONF_DIR="/etc/nginx"
            OPENRESTY_USER="nginx"
            ;;
        *)
            log_error "Unsupported distribution: $ID"
            exit 1
            ;;
    esac
else
    log_error "Cannot detect distribution. /etc/os-release not found."
    exit 1
fi

log_info "Detected distribution: $PRETTY_NAME"
log_info "Using package manager: $PKG_MGR"

# Check if OpenResty is installed
log_info "[2/8] Checking OpenResty installation..."
if ! command -v openresty >/dev/null 2>&1 && ! command -v nginx >/dev/null 2>&1; then
    log_error "OpenResty or Nginx not found. Please install OpenResty first."
    exit 1
fi

# Update system packages
log_info "[3/8] Updating system packages..."
$PKG_UPDATE

# Install required packages
log_info "[4/8] Installing required packages..."
$PKG_INSTALL curl wget git

if [[ "$PKG_MGR" == "apt" ]]; then
    $PKG_INSTALL build-essential
else
    $PKG_INSTALL gcc gcc-c++ make
fi

# Download nginx-lua-prometheus module
log_info "[5/8] Installing nginx-lua-prometheus module..."
if [[ -d /opt/nginx-lua-prometheus ]]; then
    rm -rf /opt/nginx-lua-prometheus
fi

cd /opt
git clone https://github.com/knyar/nginx-lua-prometheus.git
chown -R root:root nginx-lua-prometheus
chmod -R 755 nginx-lua-prometheus

# Create prometheus metrics configuration
log_info "[6/8] Creating Prometheus metrics configuration..."
mkdir -p "${OPENRESTY_CONF_DIR}/conf.d"

cat > "${OPENRESTY_CONF_DIR}/conf.d/prometheus-metrics.conf" << EOF
# Prometheus metrics configuration
lua_shared_dict prometheus_metrics 10M;
lua_package_path "/opt/nginx-lua-prometheus/?.lua;/usr/local/openresty/lualib/?.lua;;";

init_worker_by_lua_block {
    prometheus = require("prometheus").init("prometheus_metrics")
    
    -- HTTP request metrics
    metric_requests = prometheus:counter(
        "nginx_http_requests_total", "Number of HTTP requests", {"host", "status", "method"}
    )
    
    metric_latency = prometheus:histogram(
        "nginx_http_request_duration_seconds", "HTTP request latency", {"host"}
    )
    
    -- Connection metrics
    metric_connections = prometheus:gauge(
        "nginx_http_connections", "Number of HTTP connections", {"state"}
    )
    
    -- Upstream metrics
    metric_upstream_latency = prometheus:histogram(
        "nginx_upstream_response_time_seconds", "Upstream response time", {"upstream"}
    )
}

log_by_lua_block {
    local host = ngx.var.host or "unknown"
    local status = tostring(ngx.var.status)
    local method = ngx.var.request_method
    local request_time = tonumber(ngx.var.request_time)
    local upstream_response_time = tonumber(ngx.var.upstream_response_time)
    
    -- Record request metrics
    metric_requests:inc(1, {host, status, method})
    
    if request_time then
        metric_latency:observe(request_time, {host})
    end
    
    -- Record upstream metrics
    if upstream_response_time then
        local upstream_addr = ngx.var.upstream_addr or "unknown"
        metric_upstream_latency:observe(upstream_response_time, {upstream_addr})
    end
    
    -- Update connection metrics
    local connections_active = ngx.var.connections_active or 0
    local connections_reading = ngx.var.connections_reading or 0
    local connections_writing = ngx.var.connections_writing or 0
    local connections_waiting = ngx.var.connections_waiting or 0
    
    metric_connections:set(connections_active, {"active"})
    metric_connections:set(connections_reading, {"reading"})
    metric_connections:set(connections_writing, {"writing"})
    metric_connections:set(connections_waiting, {"waiting"})
}

# Metrics endpoint server
server {
    listen ${METRICS_PORT};
    server_name localhost;
    access_log off;
    
    location /metrics {
        content_by_lua_block {
            prometheus:collect()
        }
        allow 127.0.0.1;
        allow ::1;
        deny all;
    }
    
    location /nginx_status {
        stub_status on;
        access_log off;
        allow 127.0.0.1;
        allow ::1;
        deny all;
    }
}
EOF

chown root:root "${OPENRESTY_CONF_DIR}/conf.d/prometheus-metrics.conf"
chmod 644 "${OPENRESTY_CONF_DIR}/conf.d/prometheus-metrics.conf"

# Configure firewall
log_info "[7/8] Configuring firewall..."
if command -v firewall-cmd >/dev/null 2>&1 && systemctl is-active firewalld >/dev/null 2>&1; then
    firewall-cmd --permanent --add-port="${METRICS_PORT}/tcp" --zone=internal
    firewall-cmd --reload
    log_info "Firewall configured for port $METRICS_PORT (internal zone)"
elif command -v ufw >/dev/null 2>&1 && ufw status | grep -q "Status: active"; then
    ufw allow from 127.0.0.1 to any port "$METRICS_PORT"
    ufw allow from ::1 to any port "$METRICS_PORT"
    log_info "UFW configured for port $METRICS_PORT (localhost only)"
fi

# Test configuration and restart service
log_info "[8/8] Testing configuration and restarting OpenResty..."
if command -v openresty >/dev/null 2>&1; then
    openresty -t
    systemctl restart openresty || systemctl restart nginx
    systemctl enable openresty || systemctl enable nginx
else
    nginx -t
    systemctl restart nginx
    systemctl enable nginx
fi

# Verification checks
log_info "Verifying installation..."

# Check if service is running
if systemctl is-active openresty >/dev/null 2>&1 || systemctl is-active nginx >/dev/null 2>&1; then
    log_info "✓ OpenResty/Nginx is running"
else
    log_error "✗ OpenResty/Nginx is not running"
    exit 1
fi

# Check if metrics endpoint is accessible
sleep 2
if curl -s "http://localhost:${METRICS_PORT}/metrics" >/dev/null; then
    log_info "✓ Metrics endpoint is accessible at http://localhost:${METRICS_PORT}/metrics"
else
    log_warn "⚠ Metrics endpoint may not be ready yet. Check manually: curl http://localhost:${METRICS_PORT}/metrics"
fi

# Check if stub_status is accessible
if curl -s "http://localhost:${METRICS_PORT}/nginx_status" >/dev/null; then
    log_info "✓ Nginx status endpoint is accessible at http://localhost:${METRICS_PORT}/nginx_status"
else
    log_warn "⚠ Nginx status endpoint may not be ready yet"
fi

log_info "OpenResty Prometheus monitoring setup completed successfully!"
log_info ""
log_info "Next steps:"
log_info "1. Configure Prometheus to scrape metrics from http://localhost:${METRICS_PORT}/metrics"
log_info "2. Import Grafana dashboards for OpenResty/Nginx monitoring"
log_info "3. Generate some traffic to see metrics in action"