#!/bin/bash

# Environment Freeze Script
# Captures the current state of a Python environment for reproducibility

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

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

log() { echo -e "${GREEN}[INFO]${NC} $1"; }
warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

# Detect environment type
detect_environment() {
    if [ -n "$CONDA_DEFAULT_ENV" ]; then
        ENV_TYPE="conda"
        ENV_NAME="$CONDA_DEFAULT_ENV"
        log "Conda environment detected: $ENV_NAME"
    elif [ -n "$VIRTUAL_ENV" ]; then
        ENV_TYPE="venv"
        ENV_NAME=$(basename "$VIRTUAL_ENV")
        log "Virtual environment detected: $ENV_NAME"
    else
        error "No active Python environment detected!"
    fi
}

# Freeze conda environment
freeze_conda() {
    local output_dir="${SCRIPT_DIR}/frozen_envs/conda_${ENV_NAME}_${TIMESTAMP}"
    mkdir -p "$output_dir"
    
    log "Freezing conda environment to $output_dir"
    
    # Export environment.yml
    conda env export > "${output_dir}/environment.yml"
    log "✓ Exported environment.yml"
    
    # Export explicit package list
    conda list --explicit > "${output_dir}/conda_explicit.txt"
    log "✓ Exported explicit package list"
    
    # Export pip freeze
    pip freeze > "${output_dir}/pip_freeze.txt"
    log "✓ Exported pip freeze"
    
    # Export conda info
    conda info > "${output_dir}/conda_info.txt"
    log "✓ Exported conda info"
    
    # Create restoration script
    cat > "${output_dir}/restore_env.sh" << EOF
#!/bin/bash
# Restore conda environment from frozen state
# Created: $(date)

echo "Restoring conda environment: ${ENV_NAME}"
echo "From snapshot: ${TIMESTAMP}"

# Method 1: From environment.yml (recommended)
conda env create -f environment.yml -n ${ENV_NAME}_restored

# Method 2: From explicit list (alternative)
# conda create -n ${ENV_NAME}_restored --file conda_explicit.txt

echo "Environment restored as: ${ENV_NAME}_restored"
echo "Activate with: conda activate ${ENV_NAME}_restored"
EOF
    
    chmod +x "${output_dir}/restore_env.sh"
    log "✓ Created restoration script"
    
    echo ""
    log "Environment frozen successfully!"
    log "Files saved to: $output_dir"
}

# Freeze virtualenv
freeze_venv() {
    local output_dir="${SCRIPT_DIR}/frozen_envs/venv_${ENV_NAME}_${TIMESTAMP}"
    mkdir -p "$output_dir"
    
    log "Freezing virtual environment to $output_dir"
    
    # Export pip freeze
    pip freeze > "${output_dir}/requirements.txt"
    log "✓ Exported requirements.txt"
    
    # Export pip list with versions
    pip list --format=json > "${output_dir}/pip_list.json"
    log "✓ Exported pip list (JSON)"
    
    # Export Python version
    python --version > "${output_dir}/python_version.txt"
    log "✓ Exported Python version"
    
    # Get system packages if available
    if command -v dpkg &> /dev/null; then
        dpkg -l | grep -E "python|cuda|cudnn" > "${output_dir}/system_packages.txt" 2>/dev/null || true
        log "✓ Exported system packages"
    fi
    
    # Create restoration script
    cat > "${output_dir}/restore_env.sh" << EOF
#!/bin/bash
# Restore virtual environment from frozen state
# Created: $(date)

echo "Restoring virtual environment: ${ENV_NAME}"
echo "From snapshot: ${TIMESTAMP}"

# Check Python version
PYTHON_VERSION=\$(cat python_version.txt | cut -d' ' -f2 | cut -d'.' -f1,2)
echo "Required Python version: \$PYTHON_VERSION"

# Create new virtual environment
python\${PYTHON_VERSION} -m venv ${ENV_NAME}_restored || python3 -m venv ${ENV_NAME}_restored

# Activate and install packages
source ${ENV_NAME}_restored/bin/activate
pip install --upgrade pip setuptools wheel
pip install -r requirements.txt

echo "Environment restored as: ${ENV_NAME}_restored"
echo "Activate with: source ${ENV_NAME}_restored/bin/activate"
EOF
    
    chmod +x "${output_dir}/restore_env.sh"
    log "✓ Created restoration script"
    
    echo ""
    log "Environment frozen successfully!"
    log "Files saved to: $output_dir"
}

# Create comparison report
create_comparison() {
    local output_dir="${SCRIPT_DIR}/frozen_envs"
    local latest_freeze=$(ls -td "${output_dir}"/*_${ENV_NAME}_* 2>/dev/null | head -1)
    
    if [ -z "$latest_freeze" ] || [ "$latest_freeze" == "${output_dir}/${ENV_TYPE}_${ENV_NAME}_${TIMESTAMP}" ]; then
        log "No previous freeze to compare with"
        return
    fi
    
    log "Creating comparison with previous freeze..."
    
    local comparison_file="${output_dir}/${ENV_TYPE}_${ENV_NAME}_${TIMESTAMP}/comparison.txt"
    
    echo "Environment Comparison Report" > "$comparison_file"
    echo "=============================" >> "$comparison_file"
    echo "Current: ${TIMESTAMP}" >> "$comparison_file"
    echo "Previous: $(basename "$latest_freeze")" >> "$comparison_file"
    echo "" >> "$comparison_file"
    
    # Compare package lists
    if [ "$ENV_TYPE" == "venv" ]; then
        diff -u "${latest_freeze}/requirements.txt" \
                "${output_dir}/${ENV_TYPE}_${ENV_NAME}_${TIMESTAMP}/requirements.txt" \
                >> "$comparison_file" 2>&1 || true
    else
        diff -u "${latest_freeze}/pip_freeze.txt" \
                "${output_dir}/${ENV_TYPE}_${ENV_NAME}_${TIMESTAMP}/pip_freeze.txt" \
                >> "$comparison_file" 2>&1 || true
    fi
    
    log "✓ Created comparison report"
}

# Main execution
main() {
    detect_environment
    
    if [ "$ENV_TYPE" == "conda" ]; then
        freeze_conda
    else
        freeze_venv
    fi
    
    create_comparison
    
    # Create latest symlink
    ln -sfn "$(basename "${SCRIPT_DIR}/frozen_envs/${ENV_TYPE}_${ENV_NAME}_${TIMESTAMP}")" \
            "${SCRIPT_DIR}/frozen_envs/${ENV_TYPE}_${ENV_NAME}_latest"
    
    log "Created symlink: ${ENV_TYPE}_${ENV_NAME}_latest"
}

main "$@"