"""
ROCm Forge — Deployment Templates
Dockerfile, bash scripts, and requirements.txt templates for ROCm deployment.
"""

# =============================================================================
# Dockerfile Template for ROCm
# =============================================================================
DOCKERFILE_TEMPLATE = """# ============================================================
# ROCm Forge — Auto-Generated Dockerfile for AMD GPU Deployment
# Generated by ROCm Forge: The CUDA-to-AMD Migration Agent
# ============================================================

# Base image: ROCm {rocm_version} with Ubuntu 22.04
FROM rocm/dev-ubuntu-22.04:{rocm_version}-complete

# Set environment variables for ROCm
ENV ROCM_HOME=/opt/rocm
ENV PATH=$ROCM_HOME/bin:$ROCM_HOME/opencl/bin:$PATH
ENV LD_LIBRARY_PATH=$ROCM_HOME/lib:$ROCM_HOME/lib64:$LD_LIBRARY_PATH
ENV HIP_VISIBLE_DEVICES=0
ENV HSA_OVERRIDE_GFX_VERSION={gfx_version}
ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942"

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    python3 \\
    python3-pip \\
    python3-venv \\
    git \\
    wget \\
    && rm -rf /var/lib/apt/lists/*

# Create working directory
WORKDIR /app

# Install PyTorch for ROCm
RUN pip3 install --no-cache-dir \\
    torch --index-url https://download.pytorch.org/whl/rocm{rocm_version}

# Install ROCm-compatible dependencies
{pip_install_commands}

# Copy application code
COPY . /app/

# Verify GPU access
RUN python3 -c "import torch; print('PyTorch:', torch.__version__); print('HIP available:', torch.cuda.is_available())"

# Default command
CMD ["python3", "{entry_point}"]
"""


# =============================================================================
# Docker Compose Template
# =============================================================================
DOCKER_COMPOSE_TEMPLATE = """# ============================================================
# ROCm Forge — Auto-Generated Docker Compose for AMD GPU
# ============================================================
version: '3.8'

services:
  app:
    build: .
    devices:
      - /dev/kfd:/dev/kfd
      - /dev/dri:/dev/dri
    group_add:
      - video
      - render
    ipc: host
    shm_size: '16gb'
    environment:
      - HIP_VISIBLE_DEVICES=0
      - HSA_OVERRIDE_GFX_VERSION={gfx_version}
      - PYTORCH_ROCM_ARCH=gfx90a;gfx942
    volumes:
      - ./data:/app/data
      - ./models:/app/models
    ports:
      - "{port}:{port}"
"""


# =============================================================================
# Deployment Bash Script Template
# =============================================================================
DEPLOY_SCRIPT_TEMPLATE = """#!/usr/bin/env bash
# ============================================================
# ROCm Forge — Auto-Generated Deployment Script
# Deploy AI workload on AMD Developer Cloud / ROCm
# ============================================================
set -euo pipefail

echo "=============================================="
echo " ROCm Forge — AMD GPU Deployment Script"
echo "=============================================="
echo ""

# ----------------------------------------------------------
# Step 1: Verify ROCm Installation
# ----------------------------------------------------------
echo "[1/6] Verifying ROCm installation..."
if command -v rocminfo &> /dev/null; then
    echo "  ✅ ROCm is installed"
    rocminfo | grep -E "Name:|Marketing" | head -4
else
    echo "  ❌ ROCm not found. Install from: https://rocm.docs.amd.com/"
    echo "  Quick install: sudo apt install rocm-dev"
    exit 1
fi
echo ""

# ----------------------------------------------------------
# Step 2: Verify GPU Visibility
# ----------------------------------------------------------
echo "[2/6] Checking GPU visibility..."
if command -v rocm-smi &> /dev/null; then
    rocm-smi --showid --showproductname 2>/dev/null || echo "  ⚠️ rocm-smi warning (non-critical)"
    echo "  ✅ GPU check complete"
else
    echo "  ⚠️ rocm-smi not found, skipping GPU check"
fi
echo ""

# ----------------------------------------------------------
# Step 3: Set Up Python Environment
# ----------------------------------------------------------
echo "[3/6] Setting up Python environment..."
python3 -m venv rocm_env 2>/dev/null || python3 -m venv rocm_env
source rocm_env/bin/activate
pip install --upgrade pip setuptools wheel
echo "  ✅ Virtual environment ready"
echo ""

# ----------------------------------------------------------
# Step 4: Install ROCm-Compatible PyTorch
# ----------------------------------------------------------
echo "[4/6] Installing ROCm-compatible PyTorch..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm{rocm_version}
echo "  ✅ PyTorch (ROCm) installed"
echo ""

# ----------------------------------------------------------
# Step 5: Install Project Dependencies
# ----------------------------------------------------------
echo "[5/6] Installing project dependencies..."
{pip_install_commands}
echo "  ✅ Dependencies installed"
echo ""

# ----------------------------------------------------------
# Step 6: Validate Setup
# ----------------------------------------------------------
echo "[6/6] Validating AMD GPU + PyTorch setup..."
python3 - << 'VALIDATION'
import torch
import sys

print(f"  Python:  {{sys.version.split()[0]}}")
print(f"  PyTorch: {{torch.__version__}}")
print(f"  HIP available: {{torch.cuda.is_available()}}")

if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        print(f"  GPU {{i}}: {{torch.cuda.get_device_name(i)}}")
    
    # Quick sanity test
    x = torch.randn(1000, 1000, device='cuda')
    y = torch.matmul(x, x.T)
    print(f"  Matrix multiply test: ✅ PASSED (result shape: {{y.shape}})")
    print("")
    print("  🎉 AMD GPU setup is ready!")
else:
    print("  ❌ No HIP GPU detected.")
    print("  Check: HIP_VISIBLE_DEVICES, ROCm install, GPU permissions")
    sys.exit(1)
VALIDATION

echo ""
echo "=============================================="
echo " ✅ Deployment complete! Run your app with:"
echo "    source rocm_env/bin/activate"
echo "    python3 {entry_point}"
echo "=============================================="
"""


# =============================================================================
# Requirements.txt Template
# =============================================================================
REQUIREMENTS_TEMPLATE = """# ============================================================
# ROCm Forge — Auto-Generated Requirements for AMD GPU
# Install PyTorch FIRST using:
#   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2
# Then install the rest:
#   pip install -r requirements.txt
# ============================================================

{requirements}
"""


# =============================================================================
# AMD Developer Cloud Setup Guide
# =============================================================================
AMD_CLOUD_GUIDE = """# ============================================================
# ROCm Forge — AMD Developer Cloud Quick Start
# ============================================================

## Getting Started with AMD Developer Cloud

### 1. Sign Up
   - Visit: https://www.amd.com/en/developer/resources/developer-cloud.html
   - Create an AMD Developer account
   - Apply for AMD Developer Cloud credits

### 2. Launch a GPU Instance
   - Select an AMD Instinct MI300X instance
   - Choose Ubuntu 22.04 with ROCm pre-installed
   - Instance will have ROCm, PyTorch, and essential tools ready

### 3. Connect to Your Instance
   ```bash
   ssh -i your-key.pem user@your-instance-ip
   ```

### 4. Verify GPU Access
   ```bash
   rocminfo
   rocm-smi
   python3 -c "import torch; print(torch.cuda.is_available())"
   ```

### 5. Run Your Migrated Code
   ```bash
   # Upload your code
   scp -i your-key.pem -r ./your-project user@your-instance-ip:/home/user/

   # Or clone from GitHub
   git clone https://github.com/your-repo.git

   # Run the deployment script
   chmod +x deploy_rocm.sh
   ./deploy_rocm.sh
   ```

## Useful AMD Developer Cloud Commands
   ```bash
   rocm-smi                    # GPU status
   rocm-smi --showmeminfo vram  # VRAM usage
   rocminfo                     # Device info
   hipcc --version              # HIP compiler version
   rocprof                      # Profiling
   ```
"""


# =============================================================================
# Environment Setup Script Template
# =============================================================================
ENV_SETUP_TEMPLATE = """#!/usr/bin/env bash
# ============================================================
# ROCm Forge — Environment Variables for AMD GPU
# Source this file: source env_rocm.sh
# ============================================================

# ROCm paths
export ROCM_HOME=/opt/rocm
export ROCM_PATH=/opt/rocm
export PATH=$ROCM_HOME/bin:$PATH
export LD_LIBRARY_PATH=$ROCM_HOME/lib:$ROCM_HOME/lib64:$LD_LIBRARY_PATH

# GPU visibility (adjust as needed)
export HIP_VISIBLE_DEVICES={gpu_ids}

# Architecture target (MI300X)
export PYTORCH_ROCM_ARCH="gfx90a;gfx942"
export HSA_OVERRIDE_GFX_VERSION={gfx_version}

# Performance tuning
export PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
export HIP_FORCE_DEV_KERNARG=1
export MIOPEN_FIND_MODE=3

# Memory
export GPU_MAX_HW_QUEUES=8

# Debugging (uncomment if needed)
# export HIP_LAUNCH_BLOCKING=1
# export AMD_LOG_LEVEL=4
# export MIOPEN_ENABLE_LOGGING=1

echo "ROCm environment configured."
echo "ROCm Home: $ROCM_HOME"
echo "HIP Devices: $HIP_VISIBLE_DEVICES"
"""