#!/bin/bash
# MIESC Docker Setup Wizard
# Configures Docker environment with SmartLLM support
# (deepseek-coder + mistral models)

set -e

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

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

# Print functions
print_header() {
    echo ""
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${NC}     ${BOLD}MIESC Docker Setup Wizard${NC}                               ${CYAN}║${NC}"
    echo -e "${CYAN}║${NC}     Multi-layer Intelligent Evaluation for Smart Contracts  ${CYAN}║${NC}"
    echo -e "${CYAN}║${NC}     SmartLLM: deepseek-coder + mistral                       ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
}

print_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

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

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

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

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

# Check if command exists
command_exists() {
    command -v "$1" &> /dev/null
}

# Check Docker
check_docker() {
    print_step "Checking Docker installation..."

    if ! command_exists docker; then
        print_error "Docker is not installed."
        echo "Please install Docker from: https://docs.docker.com/get-docker/"
        exit 1
    fi

    if ! docker info &> /dev/null; then
        print_error "Docker daemon is not running."
        echo "Please start Docker Desktop or the Docker service."
        exit 1
    fi

    DOCKER_VERSION=$(docker version --format '{{.Server.Version}}' 2>/dev/null || echo "unknown")
    print_success "Docker version: $DOCKER_VERSION"
}

# Check Docker Compose
check_docker_compose() {
    print_step "Checking Docker Compose..."

    if docker compose version &> /dev/null; then
        COMPOSE_CMD="docker compose"
        COMPOSE_VERSION=$(docker compose version --short 2>/dev/null || echo "unknown")
    elif command_exists docker-compose; then
        COMPOSE_CMD="docker-compose"
        COMPOSE_VERSION=$(docker-compose version --short 2>/dev/null || echo "unknown")
    else
        print_error "Docker Compose is not installed."
        echo "Please install Docker Compose: https://docs.docker.com/compose/install/"
        exit 1
    fi

    print_success "Docker Compose version: $COMPOSE_VERSION"
}

# Check disk space
check_disk_space() {
    print_step "Checking available disk space..."

    # Required: ~10GB for models + images
    REQUIRED_GB=10

    if [[ "$OSTYPE" == "darwin"* ]]; then
        # macOS
        AVAILABLE_KB=$(df -k "$PROJECT_DIR" | tail -1 | awk '{print $4}')
    else
        # Linux
        AVAILABLE_KB=$(df -k "$PROJECT_DIR" | tail -1 | awk '{print $4}')
    fi

    AVAILABLE_GB=$((AVAILABLE_KB / 1024 / 1024))

    if [ "$AVAILABLE_GB" -lt "$REQUIRED_GB" ]; then
        print_warning "Low disk space: ${AVAILABLE_GB}GB available, ${REQUIRED_GB}GB recommended"
        echo "  - LLM models require ~8GB"
        echo "  - Docker images require ~2GB"
        read -p "Continue anyway? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    else
        print_success "Available disk space: ${AVAILABLE_GB}GB"
    fi
}

# Detect GPU
detect_gpu() {
    print_step "Detecting GPU availability..."

    GPU_AVAILABLE=false
    GPU_TYPE="none"

    # Check for NVIDIA GPU
    if command_exists nvidia-smi; then
        if nvidia-smi &> /dev/null; then
            GPU_AVAILABLE=true
            GPU_TYPE="nvidia"
            GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)
            print_success "NVIDIA GPU detected: $GPU_NAME"
        fi
    fi

    # Check for Apple Silicon (MPS)
    if [[ "$OSTYPE" == "darwin"* ]]; then
        if sysctl -n machdep.cpu.brand_string 2>/dev/null | grep -q "Apple"; then
            GPU_AVAILABLE=true
            GPU_TYPE="apple"
            print_success "Apple Silicon detected (Metal Performance Shaders)"
        fi
    fi

    if [ "$GPU_AVAILABLE" = false ]; then
        print_warning "No GPU detected. LLM inference will use CPU only."
        echo "  This will be slower but fully functional."
    fi
}

# Configure environment
configure_environment() {
    print_step "Configuring environment..."

    ENV_FILE="$PROJECT_DIR/.env"

    # Create or update .env file
    cat > "$ENV_FILE" << EOF
# MIESC Docker Environment Configuration
# Generated by docker-setup.sh on $(date)

# MIESC Version
MIESC_VERSION=4.3.7

# Ollama Configuration
OLLAMA_HOST=http://ollama:11434
MIESC_LLM_MODEL=mistral:latest
MIESC_CODE_MODEL=deepseek-coder:6.7b
MIESC_LLM_TIMEOUT=180

# GPU Configuration
GPU_AVAILABLE=$GPU_AVAILABLE
GPU_TYPE=$GPU_TYPE
EOF

    print_success "Environment configured: $ENV_FILE"
}

# Select deployment mode
select_mode() {
    print_step "Select deployment mode:"
    echo ""
    echo "  1) Development (default docker-compose.yml)"
    echo "     - Basic setup, models loaded on demand"
    echo "     - Good for testing and development"
    echo ""
    echo "  2) Production (docker-compose.prod-llm.yml)"
    echo "     - Pre-loads all models"
    echo "     - Includes health checks and resource limits"
    echo "     - Recommended for deployment"
    echo ""

    read -p "Select mode [1/2] (default: 1): " MODE_CHOICE

    case "$MODE_CHOICE" in
        2)
            COMPOSE_FILE="$PROJECT_DIR/docker-compose.prod-llm.yml"
            DEPLOYMENT_MODE="production"
            ;;
        *)
            COMPOSE_FILE="$PROJECT_DIR/docker-compose.yml"
            DEPLOYMENT_MODE="development"
            ;;
    esac

    print_success "Selected mode: $DEPLOYMENT_MODE"
}

# Build and start services
start_services() {
    print_step "Building and starting services..."
    echo ""

    cd "$PROJECT_DIR"

    if [ "$DEPLOYMENT_MODE" = "production" ]; then
        # Production mode - starts everything
        echo "Starting production deployment..."
        echo "This will download ~8GB of LLM models on first run."
        echo ""
        $COMPOSE_CMD -f "$COMPOSE_FILE" up -d --build
    else
        # Development mode - start with LLM profile
        echo "Starting development deployment with LLM support..."
        echo "This will download ~8GB of LLM models on first run."
        echo ""
        $COMPOSE_CMD --profile llm up -d --build ollama ollama-init
    fi

    print_success "Services started"
}

# Wait for models
wait_for_models() {
    print_step "Waiting for LLM models to be ready..."
    echo ""

    # Wait for ollama-init to complete
    echo "Downloading models (this may take 5-15 minutes on first run)..."

    if [ "$DEPLOYMENT_MODE" = "production" ]; then
        INIT_CONTAINER="miesc-ollama-init-prod"
    else
        INIT_CONTAINER="miesc-ollama-init"
    fi

    # Follow logs until init completes
    timeout 1800 docker logs -f "$INIT_CONTAINER" 2>/dev/null || true

    # Check if models are ready
    sleep 2
    if docker exec miesc-ollama ollama list 2>/dev/null | grep -q "deepseek-coder"; then
        print_success "deepseek-coder:6.7b ready"
    else
        print_warning "deepseek-coder model may still be downloading..."
    fi

    if docker exec miesc-ollama ollama list 2>/dev/null | grep -q "mistral"; then
        print_success "mistral:latest ready"
    else
        print_warning "mistral model may still be downloading..."
    fi
}

# Run health check
run_health_check() {
    print_step "Running health check..."

    HEALTH_SCRIPT="$PROJECT_DIR/deploy/health-check.sh"
    if [ -f "$HEALTH_SCRIPT" ]; then
        bash "$HEALTH_SCRIPT"
    else
        # Basic health check
        if curl -sf http://localhost:11434/api/tags > /dev/null 2>&1; then
            print_success "Ollama API is responding"
        else
            print_warning "Ollama API not yet responding (may still be starting)"
        fi
    fi
}

# Print usage instructions
print_usage() {
    echo ""
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${NC}                  ${BOLD}Setup Complete!${NC}                              ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "${BOLD}Quick Commands:${NC}"
    echo ""
    echo "  # Run a quick audit"
    echo "  docker exec -it miesc-shell miesc audit quick contracts/VulnerableBank.sol -o /data/results.json"
    echo ""
    echo "  # Generate premium report with LLM interpretation"
    echo "  docker exec -it miesc-shell miesc report /data/results.json -t premium --llm-interpret -o /data/report.md"
    echo ""
    echo "  # Interactive shell"
    if [ "$DEPLOYMENT_MODE" = "production" ]; then
        echo "  docker exec -it miesc-prod /bin/bash"
    else
        echo "  $COMPOSE_CMD --profile dev run miesc-shell"
    fi
    echo ""
    echo "  # Check model status"
    echo "  docker exec miesc-ollama ollama list"
    echo ""
    echo "  # View logs"
    echo "  $COMPOSE_CMD logs -f ollama"
    echo ""
    echo "  # Stop services"
    if [ "$DEPLOYMENT_MODE" = "production" ]; then
        echo "  $COMPOSE_CMD -f docker-compose.prod-llm.yml down"
    else
        echo "  $COMPOSE_CMD --profile llm down"
    fi
    echo ""
    echo -e "${GREEN}Ready for smart contract security analysis!${NC}"
}

# Main execution
main() {
    print_header

    check_docker
    check_docker_compose
    check_disk_space
    detect_gpu
    configure_environment
    select_mode

    echo ""
    read -p "Start MIESC services now? (Y/n): " -n 1 -r
    echo

    if [[ ! $REPLY =~ ^[Nn]$ ]]; then
        start_services
        wait_for_models
        run_health_check
    fi

    print_usage
}

# Run main
main "$@"
