#!/bin/bash
# Attach to existing AMD Strix Halo ROCm container for development

set -e

CONTAINER_NAME="amd-strix-halo-llama-rocm"
PROJECT_PATH="/mnt/shared/DEV/repos/d-popov.com/gogo2"
CONTAINER_PROJECT_PATH="/workspace/gogo2"

echo "=================================================="
echo "  Attaching to AMD Strix Halo ROCm Container"
echo "=================================================="
echo ""

# Check if container exists and is running
if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo "❌ Container '${CONTAINER_NAME}' is not running"
    echo ""
    echo "Available containers:"
    docker ps --format "  - {{.Names}} ({{.Status}})"
    echo ""
    echo "To start it: docker start ${CONTAINER_NAME}"
    exit 1
fi

echo "✓ Container is running"
echo ""

# Check if project is already mounted or accessible
echo "Checking project accessibility..."
if docker exec $CONTAINER_NAME test -d "$CONTAINER_PROJECT_PATH" 2>/dev/null; then
    echo "✓ Project already accessible at: $CONTAINER_PROJECT_PATH"
else
    echo "⚠️  Project not mounted in container"
    echo ""
    echo "OPTION 1: Mount project directory (requires container restart)"
    echo "  Add to docker-compose or docker run:"
    echo "    -v $PROJECT_PATH:$CONTAINER_PROJECT_PATH"
    echo ""
    echo "OPTION 2: Copy project into container"
    echo "  docker cp $PROJECT_PATH $CONTAINER_NAME:/workspace/"
    echo ""
    echo "OPTION 3: Work from host's home directory mount"
    echo "  (if accessible via /home/db/...)"
    echo ""
    read -p "Copy project to container now? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "Copying project to container..."
        docker exec $CONTAINER_NAME mkdir -p /workspace
        docker cp $PROJECT_PATH $CONTAINER_NAME:/workspace/
        echo "✓ Project copied"
    else
        echo "Skipping project copy"
    fi
fi
echo ""

# Check for Python
echo "Checking Python installation..."
if docker exec $CONTAINER_NAME which python3 &>/dev/null; then
    PYTHON_VERSION=$(docker exec $CONTAINER_NAME python3 --version)
    echo "✓ Python installed: $PYTHON_VERSION"
else
    echo "⚠️  Python not installed in container"
    echo ""
    echo "Install Python? (Fedora-based container)"
    read -p "Install Python 3.12 + pip? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "Installing Python..."
        docker exec $CONTAINER_NAME dnf install -y python3.12 python3-pip python3-devel git
        docker exec $CONTAINER_NAME ln -sf /usr/bin/python3.12 /usr/bin/python3
        docker exec $CONTAINER_NAME ln -sf /usr/bin/python3.12 /usr/bin/python
        echo "✓ Python installed"
    else
        echo "Skipping Python installation"
    fi
fi
echo ""

# Check for PyTorch ROCm
echo "Checking PyTorch..."
if docker exec $CONTAINER_NAME python3 -c "import torch" &>/dev/null; then
    TORCH_INFO=$(docker exec $CONTAINER_NAME python3 -c "import torch; print(f'{torch.__version__}, CUDA: {torch.cuda.is_available()}')")
    echo "✓ PyTorch installed: $TORCH_INFO"
else
    echo "⚠️  PyTorch not installed"
    echo ""
    echo "This container should have ROCm support built-in"
    echo "Install PyTorch with ROCm support?"
    read -p "Install PyTorch ROCm? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "Installing PyTorch with ROCm..."
        docker exec $CONTAINER_NAME pip3 install torch --index-url https://download.pytorch.org/whl/rocm6.2
        echo "✓ PyTorch installed"
    fi
fi
echo ""

echo "=================================================="
echo "  Ready to Attach!"
echo "=================================================="
echo ""
echo "Container: $CONTAINER_NAME"
echo "Project:   $CONTAINER_PROJECT_PATH"
echo ""
echo "Attaching to container shell..."
echo "(Use 'exit' or Ctrl+D to detach)"
echo ""
echo "Once inside, navigate to:"
echo "  cd $CONTAINER_PROJECT_PATH"
echo ""
echo "Install project dependencies:"
echo "  pip3 install -r requirements.txt"
echo ""
echo "Run ANNOTATE:"
echo "  python3 ANNOTATE/web/app.py"
echo ""
echo "=================================================="
echo ""

# Attach to container
docker exec -it $CONTAINER_NAME bash