#!/bin/bash

# Balanced Homelab Setup Script
# Keeps essential security (Fail2Ban, monitoring) while simplifying complexity

set -euo pipefail

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

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

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

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

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

# Simple dependency check
check_deps() {
    print_info "Checking dependencies..."
    
    local missing=()
    
    if ! command -v go &> /dev/null; then
        missing+=("go")
    fi
    
    if ! command -v zig &> /dev/null; then
        missing+=("zig")
    fi
    
    if ! command -v redis-server &> /dev/null; then
        missing+=("redis-server")
    fi
    
    if ! command -v docker &> /dev/null; then
        missing+=("docker")
    fi
    
    if [[ ${#missing[@]} -gt 0 ]]; then
        print_error "Missing dependencies: ${missing[*]}"
        echo ""
        echo "Install with:"
        echo "  macOS: brew install ${missing[*]}"
        echo "  Ubuntu: sudo apt-get install ${missing[*]}"
        exit 1
    fi
    
    print_success "Dependencies OK"
}

# Simple setup
setup_project() {
    print_info "Setting up project..."
    
    # Create essential directories
    mkdir -p ssl logs configs data monitoring
    
    # Generate simple SSL cert
    if [[ ! -f "ssl/cert.pem" ]]; then
        openssl req -x509 -newkey rsa:2048 -keyout ssl/key.pem -out ssl/cert.pem \
            -days 365 -nodes -subj "/C=US/ST=State/L=City/O=Homelab/CN=localhost" \
            -addext "subjectAltName=DNS:localhost,IP:127.0.0.1" 2>/dev/null
        print_success "SSL certificates generated"
    fi
    
    # Create balanced config
    cat > configs/config.yaml << 'EOF'
base_path: "./data/experiments"

auth:
  enabled: true
  api_keys:
    homelab_user:
      hash: "5e884898da28047151d0e56f8dc6292773603d0d6aabbdd62a11ef721d1542d8" # "password"
      admin: true
      roles: ["user", "admin"]
      permissions:
        read: true
        write: true
        delete: true

server:
  address: ":9101"
  tls:
    enabled: true
    cert_file: "./ssl/cert.pem"
    key_file: "./ssl/key.pem"

security:
  rate_limit:
    enabled: true
    requests_per_minute: 30
    burst_size: 5
  ip_whitelist:
    - "127.0.0.1"
    - "::1"
    - "192.168.0.0/16"
    - "10.0.0.0/8"
    - "172.16.0.0/12"
  failed_login_lockout:
    enabled: true
    max_attempts: 3
    lockout_duration: "15m"

redis:
  url: "redis://localhost:6379"

logging:
  level: "info"
  file: "./logs/app.log"
  audit_log: "./logs/audit.log"
  access_log: "./logs/access.log"

monitoring:
  enabled: true
  metrics_port: 9090
  health_check_interval: "30s"
EOF

    print_success "Configuration created"
}

# Simple build
build_project() {
    print_info "Building project..."
    
    # Build Go apps
    go build -o bin/api-server ./cmd/api-server
    go build -o bin/worker ./cmd/worker
    go build -o bin/tui ./cmd/tui
    
    # Build Zig CLI
    cd cli && zig build && cd ..
    
    print_success "Build completed"
}

# Setup Fail2Ban
setup_fail2ban() {
    print_info "Setting up Fail2Ban..."
    
    if ! command -v fail2ban-server &> /dev/null; then
        print_warning "Fail2Ban not installed, skipping..."
        return
    fi
    
    # Create Fail2Ban configuration
    sudo mkdir -p /etc/fail2ban/jail.d 2>/dev/null || true
    
    cat > /tmp/ml-experiments-jail.conf << 'EOF'
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 3
backend = systemd

[sshd]
enabled = true
port = ssh
logpath = /var/log/auth.log
maxretry = 3

[ml-experiments-api]
enabled = true
port = 9101
filter = ml-experiments-api
logpath = ./logs/audit.log
maxretry = 5
bantime = 7200

[ml-experiments-auth]
enabled = true
filter = ml-experiments-auth
logpath = ./logs/audit.log
maxretry = 3
bantime = 3600
EOF

    # Create filter definitions
    cat > /tmp/ml-experiments-api.conf << 'EOF'
[Definition]
failregex = ^.*<HOST>.*"status":40[13].*$
ignoreregex =
EOF

    cat > /tmp/ml-experiments-auth.conf << 'EOF'
[Definition]
failregex = ^.*"event":"failed_login".*"client_ip":"<HOST>".*$
ignoreregex =
EOF

    # Try to install configurations
    if sudo cp /tmp/ml-experiments-jail.conf /etc/fail2ban/jail.d/ 2>/dev/null; then
        sudo cp /tmp/ml-experiments-*.conf /etc/fail2ban/filter.d/ 2>/dev/null || true
        sudo systemctl restart fail2ban 2>/dev/null || true
        print_success "Fail2Ban configured"
    else
        print_warning "Could not configure Fail2Ban (requires sudo)"
    fi
    
    rm -f /tmp/ml-experiments-*.conf
}

# Setup Redis
setup_redis() {
    print_info "Setting up Redis..."
    
    if ! pgrep -f "redis-server" > /dev/null; then
        redis-server --daemonize yes --port 6379
        print_success "Redis started"
    else
        print_info "Redis already running"
    fi
}

# Create simple management script
create_manage_script() {
    cat > manage.sh << 'EOF'
#!/bin/bash

# Simple management script

case "${1:-status}" in
    "start")
        echo "Starting services..."
        redis-server --daemonize yes --port 6379 2>/dev/null || true
        ./bin/api-server -config configs/config.yaml &
        echo "Services started"
        ;;
    "stop")
        echo "Stopping services..."
        pkill -f "api-server" || true
        redis-cli shutdown 2>/dev/null || true
        echo "Services stopped"
        ;;
    "status")
        echo "=== Status ==="
        if pgrep -f "redis-server" > /dev/null; then
            echo "✅ Redis: Running"
        else
            echo "❌ Redis: Stopped"
        fi
        
        if pgrep -f "api-server" > /dev/null; then
            echo "✅ API Server: Running"
        else
            echo "❌ API Server: Stopped"
        fi
        ;;
    "logs")
        echo "=== Recent Logs ==="
        tail -20 logs/app.log 2>/dev/null || echo "No logs yet"
        ;;
    "test")
        echo "=== Testing ==="
        curl -k -s https://localhost:9101/health || echo "API server not responding"
        ;;
    *)
        echo "Usage: $0 {start|stop|status|logs|test}"
        ;;
esac
EOF
    
    chmod +x manage.sh
    print_success "Management script created"
}

# Show next steps
show_next_steps() {
    print_success "Setup completed!"
    echo ""
    echo "🎉 Setup complete!"
    echo ""
    echo "Next steps:"
    echo "  1. Start services: ./tools/manage.sh start"
    echo "  2. Check status: ./tools/manage.sh status"
    echo "  3. Test API: curl -k -H 'X-API-Key: password' https://localhost:9101/health"
    echo ""
    echo "Configuration: configs/config.yaml"
    echo "Logs: logs/app.log and logs/audit.log"
    echo ""
    print_success "Ready for homelab use!"
}

# Main setup
main() {
    echo "ML Experiment Manager - Homelab Setup"
    echo "====================================="
    echo ""
    
    check_deps
    setup_project
    build_project
    setup_redis
    create_manage_script
    show_next_steps
}

main "$@"