#!/bin/bash

# Automatic Setup Script for ML Experiment Manager
# Handles complete environment setup with security features

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"
}

detect_os() {
    if [[ "$OSTYPE" == "darwin"* ]]; then
        echo "macos"
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
        echo "linux"
    else
        echo "unknown"
    fi
}

install_go() {
    print_info "Installing Go..."
    
    local os=$(detect_os)
    local go_version="1.23.0"
    
    if [[ "$os" == "macos" ]]; then
        if command -v brew &> /dev/null; then
            brew install go
        else
            print_error "Homebrew not found. Please install Go manually."
            return 1
        fi
    elif [[ "$os" == "linux" ]]; then
        wget -q "https://go.dev/dl/go${go_version}.linux-amd64.tar.gz"
        sudo rm -rf /usr/local/go
        sudo tar -C /usr/local -xzf "go${go_version}.linux-amd64.tar.gz"
        rm "go${go_version}.linux-amd64.tar.gz"
        
        # Add to PATH
        echo 'export PATH=$PATH:/usr/local/go/bin' >> ~/.bashrc
        export PATH=$PATH:/usr/local/go/bin
    fi
    
    print_success "Go installed"
}

install_zig() {
    print_info "Installing Zig..."
    
    local os=$(detect_os)
    
    if [[ "$os" == "macos" ]]; then
        if command -v brew &> /dev/null; then
            brew install zig
        else
            print_error "Homebrew not found. Please install Zig manually."
            return 1
        fi
    elif [[ "$os" == "linux" ]]; then
        # Download Zig binary
        local zig_version="0.13.0"
        wget -q "https://ziglang.org/download/${zig_version}/zig-linux-x86_64-${zig_version}.tar.xz"
        tar -xf "zig-linux-x86_64-${zig_version}.tar.xz"
        sudo mv "zig-linux-x86_64-${zig_version}/zig" /usr/local/bin/
        rm -rf "zig-linux-x86_64-${zig_version}.tar.xz" "zig-linux-x86_64-${zig_version}"
    fi
    
    print_success "Zig installed"
}

install_docker() {
    print_info "Installing Docker..."
    
    local os=$(detect_os)
    
    if [[ "$os" == "macos" ]]; then
        if command -v brew &> /dev/null; then
            brew install --cask docker
            print_warning "Docker Desktop installed. Please start it manually."
        else
            print_error "Homebrew not found. Please install Docker manually."
            return 1
        fi
    elif [[ "$os" == "linux" ]]; then
        # Install Docker using official script
        curl -fsSL https://get.docker.com -o get-docker.sh
        sudo sh get-docker.sh
        sudo usermod -aG docker $USER
        rm get-docker.sh
        
        # Start Docker
        sudo systemctl enable docker
        sudo systemctl start docker
        
        print_success "Docker installed. You may need to log out and log back in."
    fi
}

install_redis() {
    print_info "Installing Redis..."
    
    local os=$(detect_os)
    
    if [[ "$os" == "macos" ]]; then
        if command -v brew &> /dev/null; then
            brew install redis
            brew services start redis
        else
            print_error "Homebrew not found. Please install Redis manually."
            return 1
        fi
    elif [[ "$os" == "linux" ]]; then
        sudo apt-get update
        sudo apt-get install -y redis-server
        sudo systemctl enable redis-server
        sudo systemctl start redis-server
    fi
    
    print_success "Redis installed and started"
}

install_dependencies() {
    print_info "Installing dependencies..."
    
    local os=$(detect_os)
    
    # Install basic tools
    if [[ "$os" == "macos" ]]; then
        if command -v brew &> /dev/null; then
            brew install openssl curl jq
        fi
    elif [[ "$os" == "linux" ]]; then
        sudo apt-get update
        sudo apt-get install -y openssl curl jq build-essential
    fi
    
    # Install Go tools
    if command -v go &> /dev/null; then
        go install github.com/golangci/golangci-lint/cmd/golangci-lint@latest
        go install golang.org/x/tools/cmd/goimports@latest
    fi
    
    print_success "Dependencies installed"
}

setup_project() {
    print_info "Setting up project..."
    
    # Create directories
    mkdir -p bin
    mkdir -p data
    mkdir -p logs
    mkdir -p db
    mkdir -p ssl
    mkdir -p configs
    mkdir -p scripts
    
    # Build project
    if command -v make &> /dev/null; then
        make build
        if command -v zig &> /dev/null; then
            make cli-build
        fi
    else
        print_warning "Make not found, building manually..."
        go build -o bin/worker ./cmd/worker
        go build -o bin/tui ./cmd/tui
        go build -o bin/data_manager ./cmd/data_manager
        go build -o bin/user_manager ./cmd/user_manager
        go build -o bin/api-server ./cmd/api-server
        
        if command -v zig &> /dev/null; then
            cd cli && zig build && cd ..
        fi
    fi
    
    print_success "Project setup completed"
}

setup_security() {
    print_info "Setting up security features..."
    
    # Generate SSL certificates
    if command -v openssl &> /dev/null; then
        openssl req -x509 -newkey rsa:4096 -keyout ssl/key.pem -out ssl/cert.pem \
            -days 365 -nodes -subj "/C=US/ST=State/L=City/O=Organization/CN=localhost" \
            -addext "subjectAltName=DNS:localhost,IP:127.0.0.1" 2>/dev/null || {
            print_warning "Failed to generate SSL certificates"
        }
        print_success "SSL certificates generated"
    fi
    
    # Generate secure configuration
    local redis_password=$(openssl rand -base64 32 2>/dev/null || echo "dev_redis_password_123")
    local jwt_secret=$(openssl rand -base64 64 2>/dev/null || echo "dev_jwt_secret_1234567890123456789012345678901234567890123456789012345678901234")
    
    cat > configs/security-config.yaml << EOF
base_path: "/data/ml-experiments"

auth:
  enabled: true
  api_keys:
    test_user:
      hash: "$(echo -n "dev_test_api_key_12345" | sha256sum | cut -d' ' -f1)"
      admin: true
      roles: ["data_scientist", "admin"]
      permissions:
        read: true
        write: true
        delete: true

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

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

redis:
  url: "redis://localhost:6379"
  password: "${redis_password}"

logging:
  level: "info"
  file: "logs/fetch_ml.log"
  audit_log: "logs/audit.log"
EOF

    cat > .env.dev << EOF
# Development environment variables
REDIS_PASSWORD=${redis_password}
JWT_SECRET=${jwt_secret}
GRAFANA_USER=admin
GRAFANA_PASSWORD=$(openssl rand -base64 16 2>/dev/null || echo "dev_grafana_password")
EOF

    print_success "Security configuration created"
}

test_installation() {
    print_info "Testing installation..."
    
    local tests_passed=0
    local tests_total=0
    
    # Test Go
    tests_total=$((tests_total + 1))
    if command -v go &> /dev/null; then
        print_success "Go: Installed"
        tests_passed=$((tests_passed + 1))
    else
        print_error "Go: Not found"
    fi
    
    # Test Zig
    tests_total=$((tests_total + 1))
    if command -v zig &> /dev/null; then
        print_success "Zig: Installed"
        tests_passed=$((tests_passed + 1))
    else
        print_warning "Zig: Not found (optional)"
        tests_total=$((tests_total - 1))
    fi
    
    # Test Docker
    tests_total=$((tests_total + 1))
    if command -v docker &> /dev/null; then
        print_success "Docker: Installed"
        tests_passed=$((tests_passed + 1))
    else
        print_warning "Docker: Not found (optional)"
        tests_total=$((tests_total - 1))
    fi
    
    # Test Redis
    tests_total=$((tests_total + 1))
    if command -v redis-cli &> /dev/null; then
        if redis-cli ping | grep -q "PONG"; then
            print_success "Redis: Running"
            tests_passed=$((tests_passed + 1))
        else
            print_warning "Redis: Not running"
        fi
    else
        print_warning "Redis: Not found (optional)"
        tests_total=$((tests_total - 1))
    fi
    
    # Test binaries
    if [[ -f "bin/api-server" ]]; then
        tests_total=$((tests_total + 1))
        if ./bin/api-server --help > /dev/null 2>&1; then
            print_success "API Server: Built"
            tests_passed=$((tests_passed + 1))
        else
            print_error "API Server: Build failed"
        fi
    fi
    
    if [[ $tests_total -gt 0 ]]; then
        local success_rate=$((tests_passed * 100 / tests_total))
        print_info "Tests: $tests_passed/$tests_total passed ($success_rate%)"
    fi
    
    print_success "Installation testing completed"
}

show_next_steps() {
    print_success "Automatic setup completed!"
    echo
    echo "Next Steps:"
    echo "==========="
    echo ""
    echo "1. Load environment variables:"
    echo "   source .env.dev"
    echo ""
    echo "2. Start the API server:"
    echo "   ./bin/api-server -config configs/config.yaml"
    echo ""
    echo "3. Test the Zig CLI (if installed):"
    echo "   ./cli/zig-out/bin/ml --help"
    echo ""
    echo "4. Deploy with Docker (optional):"
    echo "   make docker-run"
    echo ""
    echo "5. Docker Compose deployment:"
    echo "   docker-compose up -d"
    echo ""
    echo "Configuration Files:"
    echo "  configs/config.yaml            # Main configuration"
    echo "  configs/config_local.yaml      # Local development"
    echo "  ssl/cert.pem, ssl/key.pem     # TLS certificates"
    echo ""
    echo "Documentation:"
    echo "  docs/DEPLOYMENT.md            # Deployment guide"
    echo ""
    echo "Quick Commands:"
    echo "  make help                     # Show all commands"
    echo "  make test                     # Run tests"
    echo "  docker-compose up -d          # Start services"
    echo ""
    print_success "Ready to use ML Experiment Manager!"
}

# Main setup function
main() {
    echo "ML Experiment Manager Automatic Setup"
    echo "====================================="
    echo ""
    
    print_info "Starting automatic setup..."
    echo ""
    
    # Check and install dependencies
    if ! command -v go &> /dev/null; then
        print_info "Go not found, installing..."
        install_go
    fi
    
    if ! command -v zig &> /dev/null; then
        print_info "Zig not found, installing..."
        install_zig
    fi
    
    if ! command -v docker &> /dev/null; then
        print_info "Docker not found, installing..."
        install_docker
    fi
    
    if ! command -v redis-cli &> /dev/null; then
        print_info "Redis not found, installing..."
        install_redis
    fi
    
    # Install additional dependencies
    install_dependencies
    
    # Setup project
    setup_project
    
    # Setup security
    setup_security
    
    # Test installation
    test_installation
    
    # Show next steps
    show_next_steps
}

# Handle command line arguments
case "${1:-setup}" in
    "setup")
        main
        ;;
    "deps")
        install_dependencies
        ;;
    "test")
        test_installation
        ;;
    "help"|"-h"|"--help")
        echo "Automatic Setup Script"
        echo "Usage: $0 {setup|deps|test|help}"
        echo ""
        echo "Commands:"
        echo "  setup  - Run full automatic setup"
        echo "  deps   - Install dependencies only"
        echo "  test   - Test installation"
        echo "  help   - Show this help"
        ;;
    *)
        print_error "Unknown command: $1"
        echo "Use '$0 help' for usage information"
        exit 1
        ;;
esac