#!/bin/bash
# stats.sh — Traffic statistics module for GoTelegram
# 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() {
    # 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)

    # 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=$(date +%s -d "1 minute ago" +%Y%m%d%H%M 2>/dev/null || date +%Y%m%d%H%M)
    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)
    if [[ -f "$HISTORY_FILE" ]]; then
        local last_ts=$(tail -1 "$HISTORY_FILE" 2>/dev/null | cut -d, -f1)
        local current_minute=$((ts - (ts % 60)))

        if [[ -z "$last_ts" ]] || [[ $((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
}

# 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  # site_bytes is column 3

    local now=$(date +%s)
    local result=""

    # 1 minute rate
    local ts_1m=$((now - 60))
    local bytes_now=$(tail -1 "$HISTORY_FILE" 2>/dev/null | cut -d, -f"$col_idx")
    local bytes_1m=$(awk -F, -v ts="$ts_1m" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_1m=$((bytes_now - (bytes_1m > 0 ? bytes_1m : bytes_now)))
    [[ $diff_1m -lt 0 ]] && diff_1m=0
    local rate_1m=$((diff_1m / 60))
    local bytes_1m_fmt=$(format_bytes "$diff_1m")
    local rate_1m_fmt=$(format_rate "$rate_1m")

    # 5 minute rate
    local ts_5m=$((now - 300))
    local bytes_5m=$(awk -F, -v ts="$ts_5m" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_5m=$((bytes_now - (bytes_5m > 0 ? bytes_5m : bytes_now)))
    [[ $diff_5m -lt 0 ]] && diff_5m=0
    local rate_5m=$((diff_5m / 300))
    local bytes_5m_fmt=$(format_bytes "$diff_5m")
    local rate_5m_fmt=$(format_rate "$rate_5m")

    # 60 minute rate
    local ts_60m=$((now - 3600))
    local bytes_60m=$(awk -F, -v ts="$ts_60m" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_60m=$((bytes_now - (bytes_60m > 0 ? bytes_60m : bytes_now)))
    [[ $diff_60m -lt 0 ]] && diff_60m=0
    local rate_60m=$((diff_60m / 3600))
    local bytes_60m_fmt=$(format_bytes "$diff_60m")
    local rate_60m_fmt=$(format_rate "$rate_60m")

    # 1 day total
    local ts_1d=$((now - 86400))
    local bytes_1d=$(awk -F, -v ts="$ts_1d" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_1d=$((bytes_now - (bytes_1d > 0 ? bytes_1d : bytes_now)))
    [[ $diff_1d -lt 0 ]] && diff_1d=0
    local rate_1d=$((diff_1d > 0 ? diff_1d / 86400 : 0))
    local bytes_1d_fmt=$(format_bytes "$diff_1d")
    local rate_1d_fmt=$(format_rate "$rate_1d")

    # 7 days total
    local ts_7d=$((now - 604800))
    local bytes_7d=$(awk -F, -v ts="$ts_7d" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_7d=$((bytes_now - (bytes_7d > 0 ? bytes_7d : bytes_now)))
    [[ $diff_7d -lt 0 ]] && diff_7d=0
    local rate_7d=$((diff_7d > 0 ? diff_7d / 604800 : 0))
    local bytes_7d_fmt=$(format_bytes "$diff_7d")
    local rate_7d_fmt=$(format_rate "$rate_7d")

    # 30 days total
    local ts_30d=$((now - 2592000))
    local bytes_30d=$(awk -F, -v ts="$ts_30d" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_30d=$((bytes_now - (bytes_30d > 0 ? bytes_30d : bytes_now)))
    [[ $diff_30d -lt 0 ]] && diff_30d=0
    local rate_30d=$((diff_30d > 0 ? diff_30d / 2592000 : 0))
    local bytes_30d_fmt=$(format_bytes "$diff_30d")
    local rate_30d_fmt=$(format_rate "$rate_30d")

    # 365 days total
    local ts_365d=$((now - 31536000))
    local bytes_365d=$(awk -F, -v ts="$ts_365d" '$1 >= ts' "$HISTORY_FILE" 2>/dev/null | head -1 | cut -d, -f"$col_idx")
    local diff_365d=$((bytes_now - (bytes_365d > 0 ? bytes_365d : bytes_now)))
    [[ $diff_365d -lt 0 ]] && diff_365d=0
    local rate_365d=$((diff_365d > 0 ? diff_365d / 31536000 : 0))
    local bytes_365d_fmt=$(format_bytes "$diff_365d")
    local rate_365d_fmt=$(format_rate "$rate_365d")

    # Return as pipe-delimited format for table display
    echo "$bytes_1m_fmt|$rate_1m_fmt|$bytes_5m_fmt|$rate_5m_fmt|$bytes_60m_fmt|$rate_60m_fmt|$bytes_1d_fmt|$rate_1d_fmt|$bytes_7d_fmt|$rate_7d_fmt|$bytes_30d_fmt|$rate_30d_fmt|$bytes_365d_fmt|$rate_365d_fmt"
}

# 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

    # 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