#!/bin/bash
# stats.sh — Traffic statistics module for GoTelegram v2.5.0
# Tracks proxy (telemt port 443) and site (nginx port 8443) traffic
# Uses iptables counters + real-time snapshots + historical CSV

# Color codes (from common.sh)
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

STATS_DIR="/run/gotelegram"
HISTORY_FILE="/opt/gotelegram/stats_history.csv"
SNAPSHOTS_DIR="$STATS_DIR/snapshots"
CURRENT_SNAPSHOT="$STATS_DIR/stats_current.json"
CONFIG_FILE="/opt/gotelegram/config.json"

# Initialize stats infrastructure
stats_init() {
    if ! command -v iptables &>/dev/null; then
        log_warning "iptables не найден: установите пакет iptables или запустите установку зависимостей"
        return 1
    fi

    # Create runtime directory
    mkdir -p "$STATS_DIR" "$SNAPSHOTS_DIR" 2>/dev/null
    chmod 755 "$STATS_DIR" "$SNAPSHOTS_DIR" 2>/dev/null

    # Create iptables chain if not exists
    if ! iptables -L GOTELEGRAM_STATS -n >/dev/null 2>&1; then
        iptables -N GOTELEGRAM_STATS 2>/dev/null
    fi

    # Add chain to INPUT if not already present
    if ! iptables -C INPUT -j GOTELEGRAM_STATS 2>/dev/null; then
        iptables -I INPUT -j GOTELEGRAM_STATS 2>/dev/null
    fi

    # Add rule for proxy traffic (port 443, TCP)
    if ! iptables -C GOTELEGRAM_STATS -p tcp --dport 443 2>/dev/null; then
        iptables -A GOTELEGRAM_STATS -p tcp --dport 443 2>/dev/null
    fi

    # Add rule for site traffic (loopback, port 8443, TCP)
    if ! iptables -C GOTELEGRAM_STATS -i lo -p tcp --dport 8443 2>/dev/null; then
        iptables -A GOTELEGRAM_STATS -i lo -p tcp --dport 8443 2>/dev/null
    fi

    # Initialize CSV header if file doesn't exist
    if [[ ! -f "$HISTORY_FILE" ]]; then
        echo "epoch,proxy_bytes,site_bytes" > "$HISTORY_FILE" 2>/dev/null
    fi

    # Write initial snapshot
    stats_collect
}

# Collect current traffic statistics from iptables
stats_collect() {
    local proxy_bytes=0 proxy_pkts=0 site_bytes=0 site_pkts=0
    local ts=$(date +%s)
    local temp_file=$(mktemp)

    if ! command -v iptables &>/dev/null; then
        mkdir -p "$STATS_DIR" 2>/dev/null
        echo "{\"ts\":$ts,\"proxy_bytes\":0,\"proxy_pkts\":0,\"site_bytes\":0,\"site_pkts\":0,\"error\":\"iptables_missing\"}" > "$CURRENT_SNAPSHOT" 2>/dev/null
        rm -f "$temp_file" 2>/dev/null
        return 1
    fi

    # Parse iptables output: format is "pkts bytes target"
    # We need to extract bytes (2nd column) for each rule
    local iptables_output=$(iptables -L GOTELEGRAM_STATS -v -n -x 2>/dev/null)

    # Extract counters for port 443 (proxy)
    proxy_bytes=$(echo "$iptables_output" | grep "dpt:443" | grep -v "lo" | awk '{print $2}')
    proxy_pkts=$(echo "$iptables_output" | grep "dpt:443" | grep -v "lo" | awk '{print $1}')

    # Extract counters for port 8443 on loopback (site)
    site_bytes=$(echo "$iptables_output" | grep "dpt:8443" | awk '{print $2}')
    site_pkts=$(echo "$iptables_output" | grep "dpt:8443" | awk '{print $1}')

    # Default to 0 if not found
    proxy_bytes=${proxy_bytes:-0}
    proxy_pkts=${proxy_pkts:-0}
    site_bytes=${site_bytes:-0}
    site_pkts=${site_pkts:-0}

    # Write current snapshot as JSON
    if command -v jq &>/dev/null; then
        echo "{\"ts\":$ts,\"proxy_bytes\":$proxy_bytes,\"proxy_pkts\":$proxy_pkts,\"site_bytes\":$site_bytes,\"site_pkts\":$site_pkts}" > "$CURRENT_SNAPSHOT" 2>/dev/null
    else
        cat > "$CURRENT_SNAPSHOT" 2>/dev/null <<EOF
{"ts":$ts,"proxy_bytes":$proxy_bytes,"proxy_pkts":$proxy_pkts,"site_bytes":$site_bytes,"site_pkts":$site_pkts}
EOF
    fi

    # Save snapshot for rate calculation (one per minute)
    local minute_key
    minute_key=$(date +%Y%m%d%H%M 2>/dev/null)
    local snapshot_file="$SNAPSHOTS_DIR/snap_${minute_key}.json"
    cp "$CURRENT_SNAPSHOT" "$snapshot_file" 2>/dev/null

    # Append to history CSV (once per minute, check if last entry is fresh)
    # Auto-recreate the file with header if it was deleted — otherwise the
    # collector would silently stop writing history after any wipe (v2.4.1 fix).
    if [[ ! -f "$HISTORY_FILE" ]]; then
        mkdir -p "$(dirname "$HISTORY_FILE")" 2>/dev/null
        echo "epoch,proxy_bytes,site_bytes" > "$HISTORY_FILE" 2>/dev/null
    fi

    if [[ -f "$HISTORY_FILE" ]]; then
        local last_ts
        last_ts=$(grep -E '^[0-9]' "$HISTORY_FILE" 2>/dev/null | tail -1 | cut -d, -f1)
        last_ts="${last_ts:-0}"
        local current_minute=$((ts - (ts % 60)))

        if [[ "$last_ts" -eq 0 ]] || [[ $((current_minute - last_ts)) -ge 60 ]]; then
            echo "$current_minute,$proxy_bytes,$site_bytes" >> "$HISTORY_FILE" 2>/dev/null

            # Cleanup old entries (keep only 365 days)
            stats_cleanup_history
        fi
    fi

    rm -f "$temp_file" 2>/dev/null
}

# Read current snapshot as JSON
stats_read_current() {
    if [[ -f "$CURRENT_SNAPSHOT" ]]; then
        cat "$CURRENT_SNAPSHOT"
    else
        echo "{}"
    fi
}

# Extract value from JSON (fallback if jq not available)
json_get() {
    local json="$1"
    local key="$2"

    if command -v jq &>/dev/null; then
        echo "$json" | jq -r ".${key}" 2>/dev/null || echo "0"
    else
        echo "$json" | grep -o "\"$key\":[^,}]*" | cut -d: -f2 | tr -d ' "' || echo "0"
    fi
}

# Convert bytes to human-readable format
format_bytes() {
    local bytes=$1

    if (( bytes < 1024 )); then
        printf "%.0f B" "$bytes"
    elif (( bytes < 1024 * 1024 )); then
        printf "%.1f KB" "$(echo "scale=1; $bytes / 1024" | bc 2>/dev/null || echo "$((bytes / 1024))")"
    elif (( bytes < 1024 * 1024 * 1024 )); then
        printf "%.1f MB" "$(echo "scale=1; $bytes / 1024 / 1024" | bc 2>/dev/null || echo "$((bytes / 1024 / 1024))")"
    elif (( bytes < 1024 * 1024 * 1024 * 1024 )); then
        printf "%.1f GB" "$(echo "scale=1; $bytes / 1024 / 1024 / 1024" | bc 2>/dev/null || echo "$((bytes / 1024 / 1024 / 1024))")"
    else
        printf "%.1f TB" "$(echo "scale=1; $bytes / 1024 / 1024 / 1024 / 1024" | bc 2>/dev/null || echo "$((bytes / 1024 / 1024 / 1024 / 1024))")"
    fi
}

# Convert bytes/sec to human-readable rate
format_rate() {
    local bytes_per_sec=$1

    if (( bytes_per_sec < 1024 )); then
        printf "%.0f B/s" "$bytes_per_sec"
    elif (( bytes_per_sec < 1024 * 1024 )); then
        printf "%.1f KB/s" "$(echo "scale=1; $bytes_per_sec / 1024" | bc 2>/dev/null || echo "$((bytes_per_sec / 1024))")"
    elif (( bytes_per_sec < 1024 * 1024 * 1024 )); then
        printf "%.1f MB/s" "$(echo "scale=1; $bytes_per_sec / 1024 / 1024" | bc 2>/dev/null || echo "$((bytes_per_sec / 1024 / 1024))")"
    else
        printf "%.1f GB/s" "$(echo "scale=1; $bytes_per_sec / 1024 / 1024 / 1024" | bc 2>/dev/null || echo "$((bytes_per_sec / 1024 / 1024 / 1024))")"
    fi
}

# Safely convert value to integer (returns 0 for empty/non-numeric)
_to_int() {
    local val="${1:-0}"
    # Strip non-numeric chars, default to 0
    val="${val//[^0-9]/}"
    echo "${val:-0}"
}

# Calculate diff safely (never negative, never crashes on empty)
_safe_diff() {
    local a=$(_to_int "$1")
    local b=$(_to_int "$2")
    local d=$((a - b))
    (( d < 0 )) && d=0
    echo "$d"
}

# Calculate traffic rates and totals from history
stats_calculate_rates() {
    local traffic_type="$1"  # "proxy" or "site"
    local col_idx=2  # proxy_bytes is column 2
    [[ "$traffic_type" == "site" ]] && col_idx=3

    local now
    now=$(date +%s)

    # Get latest data line (skip header with grep -E '^[0-9]')
    local bytes_now
    bytes_now=$(_to_int "$(grep -E '^[0-9]' "$HISTORY_FILE" 2>/dev/null | tail -1 | cut -d, -f"$col_idx")")

    local periods="60 300 3600 86400 604800 2592000 31536000"
    local results=""

    for secs in $periods; do
        local target_ts=$((now - secs))
        # Find closest entry at or after target timestamp (skip header)
        local old_val
        old_val=$(_to_int "$(awk -F, -v ts="$target_ts" '$1 ~ /^[0-9]/ && $1 <= ts' "$HISTORY_FILE" 2>/dev/null | tail -1 | cut -d, -f"$col_idx")")

        local diff
        diff=$(_safe_diff "$bytes_now" "$old_val")
        local rate=$(( secs > 0 ? diff / secs : 0 ))

        local bytes_fmt rate_fmt
        bytes_fmt=$(format_bytes "$diff")
        rate_fmt=$(format_rate "$rate")

        if [ -z "$results" ]; then
            results="${bytes_fmt}|${rate_fmt}"
        else
            results="${results}|${bytes_fmt}|${rate_fmt}"
        fi
    done

    echo "$results"
}

# Main display function for traffic statistics
show_traffic_stats() {
    # Ensure stats are collected
    stats_collect

    # Get current counters
    local current_json=$(stats_read_current)
    local proxy_pkts=$(json_get "$current_json" "proxy_pkts")
    local site_pkts=$(json_get "$current_json" "site_pkts")

    # Calculate rates for proxy
    local proxy_rates=$(stats_calculate_rates "proxy")
    IFS='|' read -r p1m p1mr p5m p5mr p60m p60mr p1d p1dr p7d p7dr p30d p30dr p365d p365dr <<< "$proxy_rates"

    # Calculate rates for site
    local site_rates=$(stats_calculate_rates "site")
    IFS='|' read -r s1m s1mr s5m s5mr s60m s60mr s1d s1dr s7d s7dr s30d s30dr s365d s365dr <<< "$site_rates"

    # Display proxy stats
    {
        echo ""
        echo -e "${BLUE}  Proxy (telemt, порт 443):${NC}"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        echo -e "${BLUE}  Период     │ Входящий      │ Скорость${NC}"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        printf "   %-9s │ %14s │ %s\n" "1 мин" "$p1m" "$p1mr"
        printf "   %-9s │ %14s │ %s\n" "5 мин" "$p5m" "$p5mr"
        printf "   %-9s │ %14s │ %s\n" "60 мин" "$p60m" "$p60mr"
        printf "   %-9s │ %14s │ %s\n" "1 день" "$p1d" "$p1dr"
        printf "   %-9s │ %14s │ %s\n" "7 дней" "$p7d" "$p7dr"
        printf "   %-9s │ %14s │ %s\n" "30 дней" "$p30d" "$p30dr"
        printf "   %-9s │ %14s │ %s\n" "365 дней" "$p365d" "$p365dr"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        printf "  Пакетов: %d\n\n" "$proxy_pkts"

        echo -e "${BLUE}  Сайт (nginx, порт 8443):${NC}"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        echo -e "${BLUE}  Период     │ Входящий      │ Скорость${NC}"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        printf "   %-9s │ %14s │ %s\n" "1 мин" "$s1m" "$s1mr"
        printf "   %-9s │ %14s │ %s\n" "5 мин" "$s5m" "$s5mr"
        printf "   %-9s │ %14s │ %s\n" "60 мин" "$s60m" "$s60mr"
        printf "   %-9s │ %14s │ %s\n" "1 день" "$s1d" "$s1dr"
        printf "   %-9s │ %14s │ %s\n" "7 дней" "$s7d" "$s7dr"
        printf "   %-9s │ %14s │ %s\n" "30 дней" "$s30d" "$s30dr"
        printf "   %-9s │ %14s │ %s\n" "365 дней" "$s365d" "$s365dr"
        echo -e "${BLUE}  ─────────────────────────────────────────${NC}"
        printf "  Пакетов: %d\n" "$site_pkts"
        echo ""
    } >&2
}

# Clean up history older than 365 days
stats_cleanup_history() {
    if [[ ! -f "$HISTORY_FILE" ]]; then
        return
    fi

    local now=$(date +%s)
    local ts_365d=$((now - 31536000))
    local temp_file=$(mktemp)

    # Keep header + entries from last 365 days
    {
        head -1 "$HISTORY_FILE"
        awk -F, -v ts="$ts_365d" '$1 >= ts' "$HISTORY_FILE" | tail -n +2
    } > "$temp_file" 2>/dev/null

    mv "$temp_file" "$HISTORY_FILE" 2>/dev/null
}

# Toggle stats collection on/off
toggle_stats() {
    local current_state="false"

    # Read current state from config
    if [[ -f "$CONFIG_FILE" ]] && command -v jq &>/dev/null; then
        current_state=$(jq -r '.stats_enabled // false' "$CONFIG_FILE" 2>/dev/null)
    fi

    # Toggle
    if [[ "$current_state" == "true" ]]; then
        # Disable stats
        if [[ -f "$CONFIG_FILE" ]]; then
            if command -v jq &>/dev/null; then
                jq '.stats_enabled = false' "$CONFIG_FILE" > "${CONFIG_FILE}.tmp" 2>/dev/null
                mv "${CONFIG_FILE}.tmp" "$CONFIG_FILE"
            fi
        fi

        # Remove iptables rules
        iptables -D INPUT -j GOTELEGRAM_STATS 2>/dev/null
        iptables -F GOTELEGRAM_STATS 2>/dev/null
        iptables -X GOTELEGRAM_STATS 2>/dev/null

        # Clean up directories
        rm -rf "$STATS_DIR" 2>/dev/null

        echo "Сбор статистики ОТКЛЮЧЕН" >&2
    else
        # Enable stats
        if [[ -f "$CONFIG_FILE" ]]; then
            if command -v jq &>/dev/null; then
                jq '.stats_enabled = true' "$CONFIG_FILE" > "${CONFIG_FILE}.tmp" 2>/dev/null
                mv "${CONFIG_FILE}.tmp" "$CONFIG_FILE"
            fi
        fi

        # Initialize stats collection
        stats_init

        echo "Сбор статистики ВКЛЮЧЕН" >&2
    fi
}

# Install systemd service for stats collection
install_stats_collector() {
    local service_file="/etc/systemd/system/gotelegram-stats.service"

    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        echo "Требуется root для установки сервиса" >&2
        return 1
    fi

    if ! command -v iptables &>/dev/null; then
        log_info "Установка iptables для подсчёта трафика..."
        install_pkg "$(apt_pkg_for_cmd iptables)" || {
            echo "Не удалось установить iptables" >&2
            return 1
        }
    fi

    # Get script directory (resolve symlinks)
    local script_dir=$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")
    local lib_dir=$(dirname "$script_dir")

    # Create systemd service file
    cat > "$service_file" <<'EOF'
[Unit]
Description=GoTelegram Traffic Stats Collector
After=network.target
Wants=network-online.target

[Service]
Type=simple
User=root
ExecStart=/bin/bash -c 'source /opt/gotelegram/lib/common.sh; source /opt/gotelegram/lib/stats.sh; stats_init; while true; do stats_collect; sleep 1; done'
Restart=always
RestartSec=5
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF

    chmod 644 "$service_file"
    systemctl daemon-reload
    systemctl enable gotelegram-stats.service
    systemctl start gotelegram-stats.service

    echo "Сервис gotelegram-stats установлен и запущен" >&2
}

# Remove stats collector service
remove_stats_collector() {
    if [[ $EUID -ne 0 ]]; then
        echo "Требуется root для удаления сервиса" >&2
        return 1
    fi

    systemctl stop gotelegram-stats.service 2>/dev/null
    systemctl disable gotelegram-stats.service 2>/dev/null
    rm -f /etc/systemd/system/gotelegram-stats.service
    systemctl daemon-reload

    # Remove iptables rules
    iptables -D INPUT -j GOTELEGRAM_STATS 2>/dev/null
    iptables -F GOTELEGRAM_STATS 2>/dev/null
    iptables -X GOTELEGRAM_STATS 2>/dev/null

    # Clean up directories and files
    rm -rf "$STATS_DIR" 2>/dev/null
    rm -f "$HISTORY_FILE" 2>/dev/null

    echo "Сервис статистики удалён" >&2
}

# Export functions for external use
export -f stats_init stats_collect stats_read_current stats_calculate_rates
export -f show_traffic_stats format_bytes format_rate toggle_stats
export -f stats_cleanup_history install_stats_collector remove_stats_collector
export -f json_get