#!/bin/bash

#============================================================================
# FOXHUNT HFT TRADING SYSTEM - MONITORING SETUP SCRIPT
#============================================================================
# Sets up comprehensive monitoring infrastructure with Prometheus, Grafana,
# and AlertManager for the Foxhunt HFT trading system
#
# Usage:
#   ./monitoring_setup.sh [OPTIONS]
#
# Options:
#   --mode MODE          Setup mode: docker, native, k8s (default: docker)
#   --data-dir DIR       Data directory for monitoring (default: /opt/monitoring)
#   --grafana-password   Grafana admin password (default: generated)
#   --alert-webhook URL  Alert webhook URL for notifications
#   --skip-dashboards    Skip Grafana dashboard creation
#   --help               Show this help message
#============================================================================

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
SETUP_MODE="${SETUP_MODE:-docker}"
DATA_DIR="${DATA_DIR:-/opt/monitoring}"
GRAFANA_PASSWORD="${GRAFANA_PASSWORD:-$(openssl rand -base64 12)}"
ALERT_WEBHOOK="${ALERT_WEBHOOK:-}"
SKIP_DASHBOARDS="${SKIP_DASHBOARDS:-false}"

# Service ports
PROMETHEUS_PORT=9090
GRAFANA_PORT=3000
ALERTMANAGER_PORT=9093

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

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

show_help() {
    cat << EOF
Foxhunt HFT Trading System - Monitoring Setup

Usage: $0 [OPTIONS]

OPTIONS:
    --mode MODE          Setup mode (docker, native, k8s)
    --data-dir DIR       Data directory for monitoring data
    --grafana-password   Grafana admin password
    --alert-webhook URL  Alert webhook URL for notifications
    --skip-dashboards    Skip Grafana dashboard creation
    --help               Show this help message

SETUP MODES:
    docker              Deploy monitoring stack using Docker Compose
    native              Install monitoring tools natively on the system
    k8s                 Deploy to Kubernetes cluster

EXAMPLES:
    $0                                    # Docker deployment with defaults
    $0 --mode native --data-dir /var/monitoring
    $0 --alert-webhook https://hooks.slack.com/...

EOF
}

# Parse arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --mode)
                SETUP_MODE="$2"
                shift 2
                ;;
            --data-dir)
                DATA_DIR="$2"
                shift 2
                ;;
            --grafana-password)
                GRAFANA_PASSWORD="$2"
                shift 2
                ;;
            --alert-webhook)
                ALERT_WEBHOOK="$2"
                shift 2
                ;;
            --skip-dashboards)
                SKIP_DASHBOARDS=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

# Setup monitoring directories
setup_directories() {
    log_info "Setting up monitoring directories..."
    
    local dirs=(
        "$DATA_DIR"
        "$DATA_DIR/prometheus"
        "$DATA_DIR/prometheus/data"
        "$DATA_DIR/grafana"
        "$DATA_DIR/grafana/data"
        "$DATA_DIR/grafana/dashboards"
        "$DATA_DIR/alertmanager"
        "$DATA_DIR/alertmanager/data"
    )
    
    for dir in "${dirs[@]}"; do
        mkdir -p "$dir"
        if [[ "$SETUP_MODE" == "native" ]]; then
            chown -R 1000:1000 "$dir" 2>/dev/null || true
        fi
    done
    
    log_success "Monitoring directories created"
}

# Generate Prometheus configuration
create_prometheus_config() {
    log_info "Creating Prometheus configuration..."
    
    cat > "$DATA_DIR/prometheus/prometheus.yml" << 'EOF'
# Prometheus configuration for Foxhunt HFT Trading System
global:
  scrape_interval: 5s
  evaluation_interval: 5s
  external_labels:
    monitor: 'foxhunt-hft'
    environment: 'production'

rule_files:
  - "alert_rules.yml"

alerting:
  alertmanagers:
    - static_configs:
        - targets:
          - alertmanager:9093

scrape_configs:
  # Prometheus itself
  - job_name: 'prometheus'
    static_configs:
      - targets: ['localhost:9090']
    scrape_interval: 10s

  # Foxhunt Trading Service
  - job_name: 'foxhunt-trading'
    static_configs:
      - targets: ['localhost:9090']
    scrape_interval: 1s  # High frequency for trading metrics
    metrics_path: '/metrics'
    params:
      service: ['trading']

  # Foxhunt ML Training Service
  - job_name: 'foxhunt-ml-training'
    static_configs:
      - targets: ['localhost:9091']
    scrape_interval: 5s
    metrics_path: '/metrics'
    params:
      service: ['ml-training']

  # Foxhunt Backtesting Service
  - job_name: 'foxhunt-backtesting'
    static_configs:
      - targets: ['localhost:9092']
    scrape_interval: 10s
    metrics_path: '/metrics'
    params:
      service: ['backtesting']

  # Foxhunt TLI Service
  - job_name: 'foxhunt-tli'
    static_configs:
      - targets: ['localhost:9093']
    scrape_interval: 5s
    metrics_path: '/metrics'
    params:
      service: ['tli']

  # System metrics (node_exporter)
  - job_name: 'node-exporter'
    static_configs:
      - targets: ['localhost:9100']
    scrape_interval: 5s

  # PostgreSQL metrics
  - job_name: 'postgres-exporter'
    static_configs:
      - targets: ['localhost:9187']
    scrape_interval: 10s

  # Redis metrics
  - job_name: 'redis-exporter'
    static_configs:
      - targets: ['localhost:9121']
    scrape_interval: 10s

  # Custom health check metrics
  - job_name: 'foxhunt-health'
    static_configs:
      - targets: ['localhost:9999']
    scrape_interval: 30s
    metrics_path: '/health/prometheus'

  # GPU metrics (if available)
  - job_name: 'gpu-exporter'
    static_configs:
      - targets: ['localhost:9445']
    scrape_interval: 5s
EOF

    log_success "Prometheus configuration created"
}

# Generate Prometheus alert rules
create_alert_rules() {
    log_info "Creating Prometheus alert rules..."
    
    cat > "$DATA_DIR/prometheus/alert_rules.yml" << 'EOF'
groups:
  - name: foxhunt_trading
    rules:
      # Trading latency alerts
      - alert: TradingHighLatency
        expr: foxhunt_trading_latency_ms > 50
        for: 10s
        labels:
          severity: warning
        annotations:
          summary: "Trading latency is high"
          description: "Trading service latency is {{ $value }}ms, above the 50ms threshold"

      - alert: TradingCriticalLatency
        expr: foxhunt_trading_latency_ms > 100
        for: 5s
        labels:
          severity: critical
        annotations:
          summary: "Trading latency is critically high"
          description: "Trading service latency is {{ $value }}ms, above the 100ms critical threshold"

      # Order processing alerts
      - alert: OrderProcessingFailures
        expr: rate(foxhunt_orders_failed_total[5m]) > 0.01
        for: 30s
        labels:
          severity: warning
        annotations:
          summary: "High order failure rate"
          description: "Order failure rate is {{ $value }} failures per second"

      - alert: OrderProcessingDown
        expr: rate(foxhunt_orders_processed_total[5m]) == 0
        for: 60s
        labels:
          severity: critical
        annotations:
          summary: "No orders being processed"
          description: "Trading service appears to have stopped processing orders"

  - name: foxhunt_system
    rules:
      # Service availability alerts
      - alert: ServiceDown
        expr: up == 0
        for: 30s
        labels:
          severity: critical
        annotations:
          summary: "Service {{ $labels.job }} is down"
          description: "Service {{ $labels.job }} has been down for more than 30 seconds"

      # High CPU usage
      - alert: HighCPUUsage
        expr: 100 - (avg by(instance) (rate(node_cpu_seconds_total{mode="idle"}[5m])) * 100) > 80
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "High CPU usage on {{ $labels.instance }}"
          description: "CPU usage is {{ $value }}% on {{ $labels.instance }}"

      # High memory usage
      - alert: HighMemoryUsage
        expr: (1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes)) * 100 > 80
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "High memory usage on {{ $labels.instance }}"
          description: "Memory usage is {{ $value }}% on {{ $labels.instance }}"

      # Disk space alerts
      - alert: DiskSpaceFull
        expr: (1 - (node_filesystem_free_bytes / node_filesystem_size_bytes)) * 100 > 90
        for: 1m
        labels:
          severity: critical
        annotations:
          summary: "Disk space almost full on {{ $labels.instance }}"
          description: "Disk usage is {{ $value }}% on {{ $labels.instance }}"

  - name: foxhunt_ml
    rules:
      # ML model accuracy alerts
      - alert: MLModelLowAccuracy
        expr: foxhunt_ml_model_accuracy < 0.8
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "ML model accuracy is low"
          description: "ML model {{ $labels.model }} accuracy is {{ $value }}, below 80%"

      # GPU alerts (if GPU is used)
      - alert: GPUHighTemperature
        expr: nvidia_gpu_temperature_celsius > 80
        for: 2m
        labels:
          severity: warning
        annotations:
          summary: "GPU temperature is high"
          description: "GPU {{ $labels.gpu }} temperature is {{ $value }}°C"

      - alert: GPUHighMemoryUsage
        expr: (nvidia_gpu_memory_used_bytes / nvidia_gpu_memory_total_bytes) * 100 > 90
        for: 5m
        labels:
          severity: warning
        annotations:
          summary: "GPU memory usage is high"
          description: "GPU {{ $labels.gpu }} memory usage is {{ $value }}%"
EOF

    log_success "Prometheus alert rules created"
}

# Generate AlertManager configuration
create_alertmanager_config() {
    log_info "Creating AlertManager configuration..."
    
    local webhook_config=""
    if [[ -n "$ALERT_WEBHOOK" ]]; then
        webhook_config="
  - name: 'webhook'
    webhook_configs:
    - url: '$ALERT_WEBHOOK'
      send_resolved: true
      title: 'Foxhunt Alert: {{ .GroupLabels.alertname }}'
      text: '{{ range .Alerts }}{{ .Annotations.description }}{{ end }}'
"
    fi
    
    cat > "$DATA_DIR/alertmanager/alertmanager.yml" << EOF
# AlertManager configuration for Foxhunt HFT Trading System
global:
  smtp_smarthost: 'localhost:587'
  smtp_from: 'alertmanager@foxhunt.local'

templates:
  - '/etc/alertmanager/templates/*.tmpl'

route:
  group_by: ['alertname', 'instance', 'severity']
  group_wait: 10s
  group_interval: 5m
  repeat_interval: 12h
  receiver: 'web.hook'
  routes:
  - match:
      severity: critical
    receiver: 'critical-alerts'
    group_wait: 5s
    repeat_interval: 1h
  - match:
      alertname: TradingCriticalLatency
    receiver: 'trading-alerts'
    group_wait: 1s
    repeat_interval: 5m

receivers:
  - name: 'web.hook'
    webhook_configs:
    - url: 'http://localhost:5001/'
      send_resolved: true

  - name: 'critical-alerts'
    email_configs:
    - to: 'ops@foxhunt.local'
      subject: 'CRITICAL: Foxhunt Alert {{ .GroupLabels.alertname }}'
      body: |
        Alert: {{ .GroupLabels.alertname }}
        Summary: {{ range .Alerts }}{{ .Annotations.summary }}{{ end }}
        Description: {{ range .Alerts }}{{ .Annotations.description }}{{ end }}
        
        Severity: {{ .GroupLabels.severity }}
        Instance: {{ .GroupLabels.instance }}
        Time: {{ range .Alerts }}{{ .StartsAt }}{{ end }}
${webhook_config}

  - name: 'trading-alerts'
    email_configs:
    - to: 'trading@foxhunt.local'
      subject: 'TRADING ALERT: {{ .GroupLabels.alertname }}'
      body: |
        URGENT: Trading system alert
        {{ range .Alerts }}{{ .Annotations.description }}{{ end }}
${webhook_config}

inhibit_rules:
  - source_match:
      severity: 'critical'
    target_match:
      severity: 'warning'
    equal: ['alertname', 'instance']
EOF

    log_success "AlertManager configuration created"
}

# Create Grafana dashboards
create_grafana_dashboards() {
    if [[ "$SKIP_DASHBOARDS" == "true" ]]; then
        log_warning "Skipping Grafana dashboard creation"
        return 0
    fi
    
    log_info "Creating Grafana dashboards..."
    
    # Trading Service Dashboard
    cat > "$DATA_DIR/grafana/dashboards/foxhunt-trading.json" << 'EOF'
{
  "dashboard": {
    "id": null,
    "title": "Foxhunt Trading Service",
    "tags": ["foxhunt", "trading"],
    "style": "dark",
    "timezone": "browser",
    "panels": [
      {
        "id": 1,
        "title": "Trading Latency",
        "type": "graph",
        "targets": [
          {
            "expr": "foxhunt_trading_latency_ms",
            "legendFormat": "Latency (ms)"
          }
        ],
        "yAxes": [
          {
            "label": "Milliseconds",
            "max": 100,
            "min": 0
          }
        ],
        "alert": {
          "conditions": [
            {
              "query": {
                "queryType": "",
                "refId": "A"
              },
              "reducer": {
                "params": [],
                "type": "last"
              },
              "evaluator": {
                "params": [50],
                "type": "gt"
              }
            }
          ],
          "executionErrorState": "alerting",
          "for": "10s",
          "frequency": "1s",
          "handler": 1,
          "name": "Trading Latency Alert",
          "noDataState": "no_data",
          "notifications": []
        }
      },
      {
        "id": 2,
        "title": "Orders Per Second",
        "type": "graph",
        "targets": [
          {
            "expr": "rate(foxhunt_orders_processed_total[1m])",
            "legendFormat": "Orders/sec"
          }
        ]
      },
      {
        "id": 3,
        "title": "Order Success Rate",
        "type": "graph",
        "targets": [
          {
            "expr": "rate(foxhunt_orders_successful_total[5m]) / rate(foxhunt_orders_processed_total[5m]) * 100",
            "legendFormat": "Success Rate %"
          }
        ]
      }
    ],
    "time": {
      "from": "now-1h",
      "to": "now"
    },
    "refresh": "1s"
  }
}
EOF

    # System Overview Dashboard
    cat > "$DATA_DIR/grafana/dashboards/foxhunt-system.json" << 'EOF'
{
  "dashboard": {
    "id": null,
    "title": "Foxhunt System Overview",
    "tags": ["foxhunt", "system"],
    "style": "dark",
    "timezone": "browser",
    "panels": [
      {
        "id": 1,
        "title": "CPU Usage",
        "type": "graph",
        "targets": [
          {
            "expr": "100 - (avg(rate(node_cpu_seconds_total{mode=\"idle\"}[5m])) * 100)",
            "legendFormat": "CPU Usage %"
          }
        ]
      },
      {
        "id": 2,
        "title": "Memory Usage",
        "type": "graph",
        "targets": [
          {
            "expr": "(1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes)) * 100",
            "legendFormat": "Memory Usage %"
          }
        ]
      },
      {
        "id": 3,
        "title": "Network I/O",
        "type": "graph",
        "targets": [
          {
            "expr": "rate(node_network_receive_bytes_total[5m])",
            "legendFormat": "Receive bytes/sec"
          },
          {
            "expr": "rate(node_network_transmit_bytes_total[5m])",
            "legendFormat": "Transmit bytes/sec"
          }
        ]
      }
    ],
    "time": {
      "from": "now-1h",
      "to": "now"
    },
    "refresh": "5s"
  }
}
EOF

    log_success "Grafana dashboards created"
}

# Setup Docker monitoring stack
setup_docker_monitoring() {
    log_info "Setting up Docker monitoring stack..."
    
    cat > "$DATA_DIR/docker-compose.monitoring.yml" << EOF
version: '3.8'

services:
  prometheus:
    image: prom/prometheus:latest
    container_name: foxhunt-prometheus
    ports:
      - "$PROMETHEUS_PORT:9090"
    volumes:
      - $DATA_DIR/prometheus/prometheus.yml:/etc/prometheus/prometheus.yml
      - $DATA_DIR/prometheus/alert_rules.yml:/etc/prometheus/alert_rules.yml
      - $DATA_DIR/prometheus/data:/prometheus
    command:
      - '--config.file=/etc/prometheus/prometheus.yml'
      - '--storage.tsdb.path=/prometheus'
      - '--web.console.libraries=/etc/prometheus/console_libraries'
      - '--web.console.templates=/etc/prometheus/consoles'
      - '--storage.tsdb.retention.time=15d'
      - '--web.enable-lifecycle'
      - '--web.enable-admin-api'
    restart: unless-stopped
    networks:
      - foxhunt-monitoring

  alertmanager:
    image: prom/alertmanager:latest
    container_name: foxhunt-alertmanager
    ports:
      - "$ALERTMANAGER_PORT:9093"
    volumes:
      - $DATA_DIR/alertmanager/alertmanager.yml:/etc/alertmanager/alertmanager.yml
      - $DATA_DIR/alertmanager/data:/alertmanager
    command:
      - '--config.file=/etc/alertmanager/alertmanager.yml'
      - '--storage.path=/alertmanager'
      - '--web.external-url=http://localhost:$ALERTMANAGER_PORT'
    restart: unless-stopped
    networks:
      - foxhunt-monitoring

  grafana:
    image: grafana/grafana:latest
    container_name: foxhunt-grafana
    ports:
      - "$GRAFANA_PORT:3000"
    volumes:
      - $DATA_DIR/grafana/data:/var/lib/grafana
      - $DATA_DIR/grafana/dashboards:/var/lib/grafana/dashboards
    environment:
      - GF_SECURITY_ADMIN_PASSWORD=$GRAFANA_PASSWORD
      - GF_USERS_ALLOW_SIGN_UP=false
      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/foxhunt-trading.json
    restart: unless-stopped
    networks:
      - foxhunt-monitoring

  node-exporter:
    image: prom/node-exporter:latest
    container_name: foxhunt-node-exporter
    ports:
      - "9100:9100"
    volumes:
      - /proc:/host/proc:ro
      - /sys:/host/sys:ro
      - /:/rootfs:ro
    command:
      - '--path.procfs=/host/proc'
      - '--path.rootfs=/rootfs'
      - '--path.sysfs=/host/sys'
      - '--collector.filesystem.mount-points-exclude=^/(sys|proc|dev|host|etc)($$|/)'
    restart: unless-stopped
    networks:
      - foxhunt-monitoring

networks:
  foxhunt-monitoring:
    driver: bridge
    name: foxhunt-monitoring

volumes:
  prometheus-data:
  grafana-data:
  alertmanager-data:
EOF

    log_success "Docker monitoring configuration created"
}

# Generate monitoring summary
generate_summary() {
    cat << EOF

${GREEN}=============================================================================
FOXHUNT HFT TRADING SYSTEM - MONITORING SETUP COMPLETE
=============================================================================${NC}

${BLUE}Monitoring Stack Configuration:${NC}
  • Mode:               $SETUP_MODE
  • Data Directory:     $DATA_DIR
  • Prometheus Port:    $PROMETHEUS_PORT
  • Grafana Port:       $GRAFANA_PORT
  • AlertManager Port:  $ALERTMANAGER_PORT

${BLUE}Grafana Access:${NC}
  • URL:                http://localhost:$GRAFANA_PORT
  • Username:           admin
  • Password:           $GRAFANA_PASSWORD

${BLUE}Service URLs:${NC}
  • Prometheus:         http://localhost:$PROMETHEUS_PORT
  • AlertManager:       http://localhost:$ALERTMANAGER_PORT
  • Grafana:            http://localhost:$GRAFANA_PORT

${BLUE}Management Commands:${NC}
EOF

    case "$SETUP_MODE" in
        "docker")
            cat << EOF
  • Start monitoring:   cd $DATA_DIR && docker-compose -f docker-compose.monitoring.yml up -d
  • Stop monitoring:    cd $DATA_DIR && docker-compose -f docker-compose.monitoring.yml down
  • View logs:          cd $DATA_DIR && docker-compose -f docker-compose.monitoring.yml logs -f
  • Restart service:    cd $DATA_DIR && docker-compose -f docker-compose.monitoring.yml restart <service>
EOF
            ;;
        "native")
            cat << EOF
  • Start Prometheus:   prometheus --config.file=$DATA_DIR/prometheus/prometheus.yml
  • Start AlertManager: alertmanager --config.file=$DATA_DIR/alertmanager/alertmanager.yml
  • Start Grafana:      grafana-server --config=$DATA_DIR/grafana/grafana.ini
EOF
            ;;
    esac

    cat << EOF

${BLUE}Next Steps:${NC}
  1. Start the monitoring stack using the commands above
  2. Access Grafana and configure additional dashboards
  3. Test alert configurations
  4. Set up notification channels (email, Slack, etc.)
  5. Configure retention policies for metrics

${GREEN}Monitoring setup completed successfully!${NC}

EOF
}

# Main setup logic
main() {
    parse_args "$@"
    
    log_info "Starting Foxhunt monitoring setup..."
    log_info "Mode: $SETUP_MODE, Data Directory: $DATA_DIR"
    
    setup_directories
    create_prometheus_config
    create_alert_rules
    create_alertmanager_config
    create_grafana_dashboards
    
    case "$SETUP_MODE" in
        "docker")
            setup_docker_monitoring
            ;;
        "native")
            log_warning "Native setup requires manual installation of Prometheus, Grafana, and AlertManager"
            ;;
        "k8s")
            log_warning "Kubernetes setup not implemented yet"
            ;;
    esac
    
    generate_summary
    
    log_success "Foxhunt monitoring setup completed!"
}

# Handle interruption
trap 'log_error "Monitoring setup interrupted"; exit 1' INT TERM

# Run main function
main "$@"