#!/usr/bin/env bash
# KGraph-MCP Deployment Script
# Usage: ./deploy.sh <environment> <version>

set -euo pipefail

# Configuration
ENVIRONMENT="${1:-}"
VERSION="${2:-}"
DOCKER_REGISTRY="${DOCKER_REGISTRY:-ghcr.io/basalganglia}"
IMAGE_NAME="kgraph-mcp"

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

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

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

validate_environment() {
    case "$ENVIRONMENT" in
        dev|development)
            ENVIRONMENT="dev"
            DEPLOY_HOST="${DEV_HOST:-dev.kgraph-mcp.example.com}"
            DEPLOY_USER="${DEV_USER:-deploy}"
            COMPOSE_FILE="docker-compose.dev.yml"
            ;;
        staging)
            DEPLOY_HOST="${STAGING_HOST:-staging.kgraph-mcp.example.com}"
            DEPLOY_USER="${STAGING_USER:-deploy}"
            COMPOSE_FILE="docker-compose.staging.yml"
            ;;
        prod|production)
            ENVIRONMENT="prod"
            DEPLOY_HOST="${PROD_HOST:-api.kgraph-mcp.example.com}"
            DEPLOY_USER="${PROD_USER:-deploy}"
            COMPOSE_FILE="docker-compose.prod.yml"
            ;;
        *)
            log_error "Invalid environment: $ENVIRONMENT"
            echo "Usage: $0 <dev|staging|prod> <version>"
            exit 1
            ;;
    esac
}

check_prerequisites() {
    log_info "Checking prerequisites..."
    
    # Check if docker is installed
    if ! command -v docker &> /dev/null; then
        log_error "Docker is not installed"
        exit 1
    fi
    
    # Check if we have SSH access
    if ! ssh -o ConnectTimeout=5 "${DEPLOY_USER}@${DEPLOY_HOST}" "echo 'SSH connection successful'" &> /dev/null; then
        log_error "Cannot connect to ${DEPLOY_USER}@${DEPLOY_HOST}"
        exit 1
    fi
    
    log_info "Prerequisites check passed"
}

build_and_push_image() {
    local tag="${DOCKER_REGISTRY}/${IMAGE_NAME}:${ENVIRONMENT}-${VERSION}"
    
    log_info "Building Docker image: $tag"
    docker build \
        --build-arg ENVIRONMENT="${ENVIRONMENT}" \
        --tag "$tag" \
        --file Dockerfile \
        .
    
    log_info "Pushing image to registry..."
    docker push "$tag"
    
    # Also tag as latest for the environment
    docker tag "$tag" "${DOCKER_REGISTRY}/${IMAGE_NAME}:${ENVIRONMENT}-latest"
    docker push "${DOCKER_REGISTRY}/${IMAGE_NAME}:${ENVIRONMENT}-latest"
}

deploy_to_server() {
    local tag="${DOCKER_REGISTRY}/${IMAGE_NAME}:${ENVIRONMENT}-${VERSION}"
    
    log_info "Deploying to ${DEPLOY_HOST}..."
    
    # Copy compose file to server
    scp "deployments/${COMPOSE_FILE}" "${DEPLOY_USER}@${DEPLOY_HOST}:/tmp/"
    
    # Deploy using docker-compose
    ssh "${DEPLOY_USER}@${DEPLOY_HOST}" << EOF
        set -e
        
        # Create deployment directory
        mkdir -p ~/kgraph-mcp
        cd ~/kgraph-mcp
        
        # Move compose file
        mv /tmp/${COMPOSE_FILE} docker-compose.yml
        
        # Set environment variables
        export KGRAPH_VERSION="${VERSION}"
        export KGRAPH_IMAGE="${tag}"
        export ENVIRONMENT="${ENVIRONMENT}"
        
        # Pull new image
        docker pull "${tag}"
        
        # Stop old containers
        docker-compose down --remove-orphans || true
        
        # Start new containers
        docker-compose up -d
        
        # Wait for health check
        sleep 10
        
        # Check if service is healthy
        if docker-compose ps | grep -q "unhealthy"; then
            echo "Service is unhealthy, rolling back..."
            docker-compose down
            exit 1
        fi
        
        # Clean up old images
        docker image prune -f
        
        echo "Deployment successful!"
EOF
}

run_post_deployment_checks() {
    log_info "Running post-deployment checks..."
    
    # Wait for service to be ready
    sleep 5
    
    # Check health endpoint
    local health_url="https://${DEPLOY_HOST}/health"
    local response=$(curl -s -o /dev/null -w "%{http_code}" "$health_url" || echo "000")
    
    if [ "$response" = "200" ]; then
        log_info "Health check passed"
    else
        log_error "Health check failed (HTTP $response)"
        return 1
    fi
    
    # Run smoke tests
    if [ -f "tests/smoke/smoke_test.sh" ]; then
        log_info "Running smoke tests..."
        DEPLOY_URL="https://${DEPLOY_HOST}" bash tests/smoke/smoke_test.sh
    fi
}

rollback() {
    log_warn "Rolling back deployment..."
    
    ssh "${DEPLOY_USER}@${DEPLOY_HOST}" << EOF
        cd ~/kgraph-mcp
        
        # Restore previous version
        docker-compose down
        export KGRAPH_IMAGE="${DOCKER_REGISTRY}/${IMAGE_NAME}:${ENVIRONMENT}-previous"
        docker-compose up -d
        
        echo "Rollback completed"
EOF
}

# Main execution
main() {
    log_info "Starting deployment process..."
    log_info "Environment: $ENVIRONMENT"
    log_info "Version: $VERSION"
    
    validate_environment
    check_prerequisites
    
    # Build and push image (skip for local dev)
    if [ "$ENVIRONMENT" != "dev" ] || [ "${BUILD_IMAGE:-true}" = "true" ]; then
        build_and_push_image
    fi
    
    # Deploy to server
    if ! deploy_to_server; then
        log_error "Deployment failed"
        rollback
        exit 1
    fi
    
    # Run post-deployment checks
    if ! run_post_deployment_checks; then
        log_error "Post-deployment checks failed"
        rollback
        exit 1
    fi
    
    log_info "Deployment completed successfully! 🚀"
    log_info "Application URL: https://${DEPLOY_HOST}"
}

# Validate arguments
if [ -z "$ENVIRONMENT" ] || [ -z "$VERSION" ]; then
    log_error "Missing required arguments"
    echo "Usage: $0 <environment> <version>"
    echo "Example: $0 dev abc123"
    echo "         $0 staging v1.2.0"
    echo "         $0 prod v1.2.0"
    exit 1
fi

# Run main function
main