#!/bin/bash
# Enhanced Docker Deployment Script for LlamaAgent
# Author: Nik Jois <nikjois@llamasearch.ai>

set -e

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
ENV_FILE="$PROJECT_ROOT/.env"
COMPOSE_FILE="$PROJECT_ROOT/docker-compose.enhanced.yml"
LOG_FILE="$PROJECT_ROOT/deployment.log"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Logging functions
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1" | tee -a "$LOG_FILE"
}

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

log_debug() {
    if [[ "${DEBUG:-}" == "1" ]]; then
        echo -e "${BLUE}[DEBUG]${NC} $1" | tee -a "$LOG_FILE"
    fi
}

log_step() {
    echo -e "${PURPLE}[STEP]${NC} $1" | tee -a "$LOG_FILE"
}

# Cleanup function
cleanup() {
    log "Performing cleanup..."
    
    # Stop any running containers if deployment failed
    if [[ "${DEPLOYMENT_FAILED:-}" == "1" ]]; then
        log_warn "Deployment failed, stopping containers..."
        docker-compose -f "$COMPOSE_FILE" down --remove-orphans 2>/dev/null || true
    fi
    
    # Remove temporary files
    rm -f /tmp/llamaagent_deploy_*.tmp 2>/dev/null || true
    
    log "Cleanup completed"
}

# Trap signals for cleanup
trap cleanup EXIT INT TERM

# Display header
show_header() {
    echo -e "${CYAN}"
    echo "=================================================="
    echo "    LlamaAgent Enhanced Docker Deployment"
    echo "    Author: Nik Jois <nikjois@llamasearch.ai>"
    echo "    Version: 2.0.0"
    echo "=================================================="
    echo -e "${NC}"
}

# Check prerequisites
check_prerequisites() {
    log_step "Checking prerequisites..."
    
    local missing_deps=()
    
    # Check Docker
    if ! command -v docker &> /dev/null; then
        missing_deps+=("docker")
    else
        local docker_version=$(docker --version | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
        log_info "Docker version: $docker_version"
    fi
    
    # Check Docker Compose
    if ! command -v docker-compose &> /dev/null; then
        missing_deps+=("docker-compose")
    else
        local compose_version=$(docker-compose --version | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
        log_info "Docker Compose version: $compose_version"
    fi
    
    # Check required tools
    for tool in curl jq; do
        if ! command -v "$tool" &> /dev/null; then
            missing_deps+=("$tool")
        fi
    done
    
    if [[ ${#missing_deps[@]} -gt 0 ]]; then
        log_error "Missing required dependencies: ${missing_deps[*]}"
        echo "Please install the missing dependencies and try again."
        exit 1
    fi
    
    # Check Docker daemon
    if ! docker info &> /dev/null; then
        log_error "Docker daemon is not running"
        exit 1
    fi
    
    # Check available disk space (minimum 10GB)
    local available_space=$(df . | tail -1 | awk '{print $4}')
    local min_space=$((10 * 1024 * 1024)) # 10GB in KB
    
    if [[ $available_space -lt $min_space ]]; then
        log_warn "Available disk space is less than 10GB. Deployment may fail."
    fi
    
    log_info "Prerequisites check completed"
}

# Create environment file
create_environment_file() {
    log_step "Setting up environment configuration..."
    
    if [[ ! -f "$ENV_FILE" ]]; then
        log_info "Creating environment file..."
        
        cat > "$ENV_FILE" << 'EOF'
# LlamaAgent Enhanced Production Environment Configuration
# Auto-generated by deployment script

# =============================================================================
# Build Configuration
# =============================================================================
BUILD_DATE=$(date -u +'%Y-%m-%dT%H:%M:%SZ')
VCS_REF=main
VERSION=2.0.0

# =============================================================================
# Application Configuration
# =============================================================================
APP_PORT=8000
METRICS_PORT=8001
ADMIN_PORT=8002
APP_DOMAIN=localhost
LOG_LEVEL=INFO
WORKERS=4
DEBUG=false
MAX_MEMORY=2G

# =============================================================================
# Security Configuration
# =============================================================================
SECRET_KEY=$(openssl rand -hex 32)
JWT_SECRET_KEY=$(openssl rand -hex 32)
ALLOWED_HOSTS=*
CORS_ORIGINS=*

# =============================================================================
# Database Configuration
# =============================================================================
POSTGRES_DB=llamaagent
POSTGRES_USER=llamaagent
POSTGRES_PASSWORD=$(openssl rand -base64 32 | tr -d "=+/" | cut -c1-25)
POSTGRES_PORT=5432
POSTGRES_MEMORY_LIMIT=1G
POSTGRES_MEMORY_RESERVATION=512M

# =============================================================================
# Redis Configuration
# =============================================================================
REDIS_PORT=6379
REDIS_MAX_MEMORY=512mb
REDIS_MEMORY_LIMIT=1G
REDIS_MEMORY_RESERVATION=256M

# =============================================================================
# Vector Database Configuration
# =============================================================================
QDRANT_HTTP_PORT=6333
QDRANT_GRPC_PORT=6334
QDRANT_LOG_LEVEL=INFO
QDRANT_CLUSTER_ENABLED=false
QDRANT_MEMORY_LIMIT=2G
QDRANT_MEMORY_RESERVATION=512M

# =============================================================================
# LLM Provider API Keys (Replace with your actual keys)
# =============================================================================
OPENAI_API_KEY=sk-placeholder-replace-with-your-key
ANTHROPIC_API_KEY=sk-ant-placeholder-replace-with-your-key
TOGETHER_API_KEY=placeholder-replace-with-your-key
COHERE_API_KEY=placeholder-replace-with-your-key
HUGGINGFACE_API_KEY=hf_placeholder-replace-with-your-token

# =============================================================================
# Monitoring Configuration
# =============================================================================
PROMETHEUS_ENABLED=true
PROMETHEUS_PORT=9090
GRAFANA_ENABLED=true
GRAFANA_PORT=3000
GRAFANA_USER=admin
GRAFANA_PASSWORD=$(openssl rand -base64 16 | tr -d "=+/" | cut -c1-12)
JAEGER_ENABLED=true
JAEGER_UI_PORT=16686

# =============================================================================
# Performance Configuration
# =============================================================================
CACHE_ENABLED=true
CACHE_TTL=3600
RATE_LIMIT_ENABLED=true
RATE_LIMIT_PER_MINUTE=100
CPU_LIMIT=2.0
MEMORY_LIMIT=4G
CPU_RESERVATION=1.0
MEMORY_RESERVATION=2G

# =============================================================================
# Web Server Configuration
# =============================================================================
HTTP_PORT=80
HTTPS_PORT=443

# =============================================================================
# Logging Configuration
# =============================================================================
ELASTICSEARCH_PORT=9200
ELASTICSEARCH_HEAP_SIZE=1g
KIBANA_PORT=5601

# =============================================================================
# Local LLM Configuration
# =============================================================================
OLLAMA_PORT=11434
OLLAMA_MODELS=llama3.2:3b

# =============================================================================
# Backup Configuration
# =============================================================================
BACKUP_SCHEDULE=0 2 * * *
EOF
        
        log_info "Environment file created at $ENV_FILE"
        log_warn "Please edit $ENV_FILE and update the API keys before continuing"
    else
        log_info "Environment file already exists"
    fi
}

# Validate environment configuration
validate_environment() {
    log_step "Validating environment configuration..."
    
    if [[ ! -f "$ENV_FILE" ]]; then
        log_error "Environment file not found: $ENV_FILE"
        exit 1
    fi
    
    # Source environment file
    set -a
    source "$ENV_FILE"
    set +a
    
    # Check for placeholder values
    local placeholder_keys=()
    
    if [[ "$OPENAI_API_KEY" == "sk-placeholder-replace-with-your-key" ]]; then
        placeholder_keys+=("OPENAI_API_KEY")
    fi
    
    if [[ "$ANTHROPIC_API_KEY" == "sk-ant-placeholder-replace-with-your-key" ]]; then
        placeholder_keys+=("ANTHROPIC_API_KEY")
    fi
    
    if [[ ${#placeholder_keys[@]} -gt 0 ]]; then
        log_warn "Found placeholder values for: ${placeholder_keys[*]}"
        log_warn "The system will use mock responses for these providers"
    fi
    
    # Validate required ports are available
    local ports_to_check=(8000 5432 6379 6333 9090 3000)
    local busy_ports=()
    
    for port in "${ports_to_check[@]}"; do
        if netstat -ln 2>/dev/null | grep -q ":$port "; then
            busy_ports+=("$port")
        fi
    done
    
    if [[ ${#busy_ports[@]} -gt 0 ]]; then
        log_error "The following ports are already in use: ${busy_ports[*]}"
        log_error "Please stop services using these ports or modify the configuration"
        exit 1
    fi
    
    log_info "Environment validation completed"
}

# Build Docker images
build_images() {
    log_step "Building Docker images..."
    
    cd "$PROJECT_ROOT"
    
    # Build with build args
    local build_args=(
        "--build-arg" "BUILD_DATE=$(date -u +'%Y-%m-%dT%H:%M:%SZ')"
        "--build-arg" "VCS_REF=$(git rev-parse --short HEAD 2>/dev/null || echo 'unknown')"
        "--build-arg" "VERSION=${VERSION:-2.0.0}"
    )
    
    log_info "Building enhanced production image..."
    if ! docker build "${build_args[@]}" -f Dockerfile.enhanced -t llamaagent:enhanced-latest .; then
        log_error "Failed to build Docker image"
        exit 1
    fi
    
    log_info "Docker images built successfully"
}

# Start services
start_services() {
    log_step "Starting services..."
    
    cd "$PROJECT_ROOT"
    
    # Create necessary directories
    mkdir -p logs data cache uploads exports backups monitoring
    
    # Start services in correct order
    log_info "Starting database services..."
    docker-compose -f "$COMPOSE_FILE" up -d postgres redis qdrant
    
    # Wait for databases to be ready
    log_info "Waiting for databases to be ready..."
    sleep 30
    
    # Start main application services
    log_info "Starting application services..."
    docker-compose -f "$COMPOSE_FILE" up -d llamaagent-app llamaagent-worker
    
    # Wait for application to be ready
    log_info "Waiting for application to be ready..."
    sleep 20
    
    # Start monitoring services
    log_info "Starting monitoring services..."
    docker-compose -f "$COMPOSE_FILE" up -d prometheus grafana jaeger
    
    # Start logging services
    log_info "Starting logging services..."
    docker-compose -f "$COMPOSE_FILE" up -d elasticsearch logstash kibana
    
    # Start remaining services
    log_info "Starting remaining services..."
    docker-compose -f "$COMPOSE_FILE" up -d nginx ollama health-check
    
    log_info "All services started successfully"
}

# Verify deployment
verify_deployment() {
    log_step "Verifying deployment..."
    
    local endpoints=(
        "http://localhost:8000/health|Main Application"
        "http://localhost:8000/metrics|Metrics Endpoint"
        "http://localhost:9090/-/healthy|Prometheus"
        "http://localhost:3000/api/health|Grafana"
        "http://localhost:5601/api/status|Kibana"
        "http://localhost:16686/|Jaeger"
    )
    
    local failed_checks=()
    
    for endpoint_info in "${endpoints[@]}"; do
        local endpoint=$(echo "$endpoint_info" | cut -d'|' -f1)
        local name=$(echo "$endpoint_info" | cut -d'|' -f2)
        
        log_info "Checking $name..."
        if curl -sf --max-time 10 "$endpoint" > /dev/null 2>&1; then
            log_info "✓ $name is healthy"
        else
            log_warn "✗ $name is not responding"
            failed_checks+=("$name")
        fi
    done
    
    # Check container status
    log_info "Checking container status..."
    local unhealthy_containers=$(docker-compose -f "$COMPOSE_FILE" ps --filter "status=exited" --format "table {{.Service}}")
    
    if [[ -n "$unhealthy_containers" && "$unhealthy_containers" != "SERVICE" ]]; then
        log_warn "Some containers are not running:"
        echo "$unhealthy_containers"
        failed_checks+=("Container Status")
    fi
    
    if [[ ${#failed_checks[@]} -gt 0 ]]; then
        log_warn "Some health checks failed: ${failed_checks[*]}"
        log_warn "The system may still be starting up. Please check again in a few minutes."
    else
        log_info "All health checks passed!"
    fi
}

# Display access information
show_access_info() {
    log_step "Deployment completed successfully!"
    
    echo -e "${GREEN}"
    echo "=================================================="
    echo "    LlamaAgent is now running!"
    echo "=================================================="
    echo -e "${NC}"
    
    echo "Service URLs:"
    echo "  • Main Application:     http://localhost:8000"
    echo "  • API Documentation:    http://localhost:8000/docs"
    echo "  • Metrics Dashboard:    http://localhost:8000/metrics"
    echo "  • Admin Panel:          http://localhost:8000/admin"
    echo ""
    echo "Monitoring & Analytics:"
    echo "  • Prometheus:           http://localhost:9090"
    echo "  • Grafana:              http://localhost:3000"
    echo "  • Jaeger Tracing:       http://localhost:16686"
    echo "  • Kibana Logs:          http://localhost:5601"
    echo ""
    echo "Databases:"
    echo "  • PostgreSQL:           localhost:5432"
    echo "  • Redis:                localhost:6379"
    echo "  • Qdrant Vector DB:     http://localhost:6333"
    echo ""
    echo "Local LLM:"
    echo "  • Ollama:               http://localhost:11434"
    echo ""
    
    # Display login credentials
    source "$ENV_FILE"
    echo "Default Credentials:"
    echo "  • Database User:        $POSTGRES_USER"
    echo "  • Database Password:    $POSTGRES_PASSWORD"
    echo "  • Grafana User:         $GRAFANA_USER"
    echo "  • Grafana Password:     $GRAFANA_PASSWORD"
    echo ""
    
    echo "Useful Commands:"
    echo "  • View logs:            docker-compose -f $COMPOSE_FILE logs -f"
    echo "  • Restart services:     docker-compose -f $COMPOSE_FILE restart"
    echo "  • Stop services:        docker-compose -f $COMPOSE_FILE down"
    echo "  • Update services:      $0 --update"
    echo ""
    
    echo -e "${YELLOW}Important Notes:${NC}"
    echo "  • Update API keys in $ENV_FILE for production use"
    echo "  • Configure SSL certificates for HTTPS access"
    echo "  • Set up regular backups for production data"
    echo "  • Monitor resource usage and scale as needed"
    echo ""
}

# Update deployment
update_deployment() {
    log_step "Updating deployment..."
    
    # Pull latest images
    log_info "Pulling latest images..."
    docker-compose -f "$COMPOSE_FILE" pull
    
    # Rebuild custom images
    build_images
    
    # Restart services with zero downtime
    log_info "Performing rolling update..."
    docker-compose -f "$COMPOSE_FILE" up -d --force-recreate --no-deps llamaagent-app llamaagent-worker
    
    # Verify update
    verify_deployment
    
    log_info "Update completed successfully"
}

# Main function
main() {
    local command="${1:-deploy}"
    
    show_header
    
    case "$command" in
        "deploy")
            check_prerequisites
            create_environment_file
            validate_environment
            build_images
            start_services
            verify_deployment
            show_access_info
            ;;
        "update"|"--update")
            update_deployment
            ;;
        "status"|"--status")
            verify_deployment
            ;;
        "logs"|"--logs")
            docker-compose -f "$COMPOSE_FILE" logs -f
            ;;
        "stop"|"--stop")
            log_info "Stopping services..."
            docker-compose -f "$COMPOSE_FILE" down
            ;;
        "restart"|"--restart")
            log_info "Restarting services..."
            docker-compose -f "$COMPOSE_FILE" restart
            ;;
        "--help"|"-h")
            echo "Usage: $0 [command]"
            echo ""
            echo "Commands:"
            echo "  deploy    Deploy the complete system (default)"
            echo "  update    Update the deployment"
            echo "  status    Check deployment status"
            echo "  logs      Follow logs"
            echo "  stop      Stop all services"
            echo "  restart   Restart all services"
            echo "  --help    Show this help message"
            ;;
        *)
            log_error "Unknown command: $command"
            echo "Use '$0 --help' for usage information"
            exit 1
            ;;
    esac
}

# Check if running as script (not sourced)
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi 