#!/usr/bin/env bash
set -e

# ============================================================
#  Anima base LoRA Trainer - Adaptive RunPod Installer
#  Repo: https://github.com/aitrepreneur/citron-anima-lora-trainer-ui
# ============================================================

REPO_URL="https://github.com/aitrepreneur/citron-anima-lora-trainer-ui.git"
APP_DIR="/workspace/citron-anima-lora-trainer-ui"
PORT="7860"

LINUX_MIN_CU128="570.26"
LINUX_MIN_CU126="560.35.05"
LINUX_MIN_CU118="520.61.05"

TORCH_CU128="pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cu128"
TORCH_CU126="pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cu126"
TORCH_CU118="pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu118"

echo "============================================================"
echo " Anima base LoRA Trainer - Adaptive RunPod Installer"
echo "============================================================"
echo ""
echo "This installer detects the RunPod NVIDIA GPU and driver, then chooses"
echo "the safest PyTorch CUDA wheel automatically."
echo ""
echo "Important: PyTorch brings its own CUDA runtime. The NVIDIA driver matters most."
echo ""

# ------------------------------------------------------------
# 1. System packages
# ------------------------------------------------------------
echo "[1/10] Installing system packages..."
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y \
  git \
  wget \
  curl \
  python3 \
  python3-venv \
  python3-dev \
  build-essential

# ------------------------------------------------------------
# 2. Choose Python
# ------------------------------------------------------------
echo ""
echo "[2/10] Checking Python version..."
PYTHON_BIN="python3"
python3 - <<'PYCODE'
import sys
if sys.version_info < (3, 10):
    raise SystemExit("Python 3.10 or newer is required.")
print("Python:", sys.version)
PYCODE

# ------------------------------------------------------------
# 3. Detect GPU and choose PyTorch CUDA version
# ------------------------------------------------------------
echo ""
echo "[3/10] Detecting NVIDIA GPU and choosing PyTorch CUDA version..."
cat > /tmp/detect_anima_gpu.py <<'PYCODE'
import subprocess
import re
import os
import sys

LINUX_MIN_CU128 = os.environ.get("LINUX_MIN_CU128", "570.26")
LINUX_MIN_CU126 = os.environ.get("LINUX_MIN_CU126", "560.35.05")
LINUX_MIN_CU118 = os.environ.get("LINUX_MIN_CU118", "520.61.05")

def version_tuple(value):
    nums = re.findall(r"\d+", str(value))
    nums = [int(x) for x in nums[:4]]
    while len(nums) < 4:
        nums.append(0)
    return tuple(nums)

def ge(value, minimum):
    return version_tuple(value) >= version_tuple(minimum)

def out(key, value):
    value = str(value).replace('"', '').replace('`', '')
    print(f'export {key}="{value}"')

try:
    query = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"],
        text=True,
        stderr=subprocess.STDOUT,
        encoding="utf-8",
        errors="ignore",
    ).strip().splitlines()[0]
    parts = [p.strip() for p in query.split(',')]
    gpu_name = parts[0]
    driver_version = parts[1] if len(parts) > 1 else "0"
    full_smi = subprocess.check_output(
        ["nvidia-smi"], text=True, stderr=subprocess.STDOUT, encoding="utf-8", errors="ignore"
    )
except Exception as e:
    out("ANIMA_GPU_OK", "0")
    out("ANIMA_GPU_NAME", "NONE")
    out("ANIMA_DRIVER_VERSION", "0")
    out("ANIMA_NVIDIA_SMI_CUDA", "Unknown")
    out("ANIMA_TORCH_CUDA", "NONE")
    out("ANIMA_MIXED_PRECISION", "fp16")
    out("ANIMA_GPU_REASON", f"nvidia-smi failed: {e}")
    sys.exit(0)

m = re.search(r"CUDA Version:\s*([0-9.]+)", full_smi)
nvidia_smi_cuda = m.group(1) if m else "Unknown"
name = gpu_name.lower()

is_blackwell = bool(re.search(r"rtx\s*50|rtx\s*5090|rtx\s*5080|rtx\s*5070|rtx\s*5060|\b5090\b|\b5080\b|\b5070\b|\b5060\b|blackwell", name))
is_pascal_or_maxwell = bool(re.search(r"gtx\s*10|gtx\s*9|\b1050\b|\b1060\b|\b1070\b|\b1080\b|\b950\b|\b960\b|\b970\b|\b980\b|titan xp|titan x\b|quadro\s*p|tesla\s*p|p100|p40|p4|maxwell|pascal", name))
is_kepler_or_older = bool(re.search(r"gtx\s*[45678]\d\d|\b750\b|\b760\b|\b770\b|\b780\b|kepler|fermi", name))
is_turing = bool(re.search(r"rtx\s*20|\b2060\b|\b2070\b|\b2080\b|gtx\s*16|\b1650\b|\b1660\b|t4|turing", name))
is_ampere_or_newer = bool(re.search(r"rtx\s*30|rtx\s*40|rtx\s*50|\b3050\b|\b3060\b|\b3070\b|\b3080\b|\b3090\b|\b4050\b|\b4060\b|\b4070\b|\b4080\b|\b4090\b|a100|a10|a16|a30|a40|a4000|a4500|a5000|a6000|l4|l40|h100|h200|b100|b200|blackwell|ampere|ada|hopper", name))

if is_kepler_or_older:
    out("ANIMA_GPU_OK", "0")
    out("ANIMA_GPU_NAME", gpu_name)
    out("ANIMA_DRIVER_VERSION", driver_version)
    out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
    out("ANIMA_TORCH_CUDA", "UNSUPPORTED")
    out("ANIMA_MIXED_PRECISION", "fp16")
    out("ANIMA_GPU_REASON", "GPU looks too old for this trainer. Recommended minimum is GTX 10-series 6GB, RTX 20-series, or newer.")
    sys.exit(0)

if is_blackwell:
    if not ge(driver_version, LINUX_MIN_CU128):
        out("ANIMA_GPU_OK", "0")
        out("ANIMA_GPU_NAME", gpu_name)
        out("ANIMA_DRIVER_VERSION", driver_version)
        out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
        out("ANIMA_TORCH_CUDA", "cu128")
        out("ANIMA_MIXED_PRECISION", "bf16")
        out("ANIMA_GPU_REASON", f"RTX 50-series detected, but driver is too old. Use NVIDIA driver {LINUX_MIN_CU128} or newer.")
        sys.exit(0)
    out("ANIMA_GPU_OK", "1")
    out("ANIMA_GPU_NAME", gpu_name)
    out("ANIMA_DRIVER_VERSION", driver_version)
    out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
    out("ANIMA_TORCH_CUDA", "cu128")
    out("ANIMA_MIXED_PRECISION", "bf16")
    out("ANIMA_GPU_REASON", "RTX 50-series / Blackwell detected. Using PyTorch CUDA 12.8 for sm_120 support.")
    sys.exit(0)

if is_pascal_or_maxwell:
    mixed = "fp16"
    if ge(driver_version, LINUX_MIN_CU126):
        torch_cuda = "cu126"
        reason = "Pascal/Maxwell legacy GPU detected. Avoiding CUDA 12.8 because modern cu128 PyTorch wheels may drop Pascal/Maxwell support. Using CUDA 12.6."
    elif ge(driver_version, LINUX_MIN_CU118):
        torch_cuda = "cu118"
        reason = "Pascal/Maxwell legacy GPU detected with older driver. Using CUDA 11.8 fallback."
    else:
        out("ANIMA_GPU_OK", "0")
        out("ANIMA_GPU_NAME", gpu_name)
        out("ANIMA_DRIVER_VERSION", driver_version)
        out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
        out("ANIMA_TORCH_CUDA", "cu118")
        out("ANIMA_MIXED_PRECISION", mixed)
        out("ANIMA_GPU_REASON", f"Driver is too old. Use NVIDIA driver {LINUX_MIN_CU118} or newer.")
        sys.exit(0)
    out("ANIMA_GPU_OK", "1")
    out("ANIMA_GPU_NAME", gpu_name)
    out("ANIMA_DRIVER_VERSION", driver_version)
    out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
    out("ANIMA_TORCH_CUDA", torch_cuda)
    out("ANIMA_MIXED_PRECISION", mixed)
    out("ANIMA_GPU_REASON", reason)
    sys.exit(0)

if ge(driver_version, LINUX_MIN_CU128):
    torch_cuda = "cu128"
    reason = "Modern NVIDIA GPU and new driver detected. Using PyTorch CUDA 12.8."
elif ge(driver_version, LINUX_MIN_CU126):
    torch_cuda = "cu126"
    reason = "Modern NVIDIA GPU detected, but driver is below CUDA 12.8 minimum. Using CUDA 12.6 fallback."
elif ge(driver_version, LINUX_MIN_CU118):
    torch_cuda = "cu118"
    reason = "Modern NVIDIA GPU detected with older driver. Using CUDA 11.8 fallback."
else:
    out("ANIMA_GPU_OK", "0")
    out("ANIMA_GPU_NAME", gpu_name)
    out("ANIMA_DRIVER_VERSION", driver_version)
    out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
    out("ANIMA_TORCH_CUDA", "cu118")
    out("ANIMA_MIXED_PRECISION", "fp16")
    out("ANIMA_GPU_REASON", f"Driver is too old. Use NVIDIA driver {LINUX_MIN_CU118} or newer.")
    sys.exit(0)

mixed = "bf16" if is_ampere_or_newer else "fp16"
if is_turing:
    mixed = "fp16"

out("ANIMA_GPU_OK", "1")
out("ANIMA_GPU_NAME", gpu_name)
out("ANIMA_DRIVER_VERSION", driver_version)
out("ANIMA_NVIDIA_SMI_CUDA", nvidia_smi_cuda)
out("ANIMA_TORCH_CUDA", torch_cuda)
out("ANIMA_MIXED_PRECISION", mixed)
out("ANIMA_GPU_REASON", reason)
PYCODE

export LINUX_MIN_CU128 LINUX_MIN_CU126 LINUX_MIN_CU118
python3 /tmp/detect_anima_gpu.py > /tmp/anima_gpu_choice.sh
source /tmp/anima_gpu_choice.sh

echo "GPU: ${ANIMA_GPU_NAME}"
echo "Driver: ${ANIMA_DRIVER_VERSION}"
echo "nvidia-smi CUDA shown: ${ANIMA_NVIDIA_SMI_CUDA}"
echo "Selected PyTorch CUDA: ${ANIMA_TORCH_CUDA}"
echo "Mixed precision default: ${ANIMA_MIXED_PRECISION}"
echo "Reason: ${ANIMA_GPU_REASON}"
echo ""

if [ "${ANIMA_GPU_OK}" != "1" ]; then
  echo "ERROR: This GPU/driver setup is not supported by the automatic installer."
  echo "${ANIMA_GPU_REASON}"
  exit 1
fi

case "${ANIMA_TORCH_CUDA}" in
  cu128) TORCH_INSTALL_CMD="${TORCH_CU128}" ;;
  cu126) TORCH_INSTALL_CMD="${TORCH_CU126}" ;;
  cu118) TORCH_INSTALL_CMD="${TORCH_CU118}" ;;
  *) echo "ERROR: Unknown PyTorch CUDA choice: ${ANIMA_TORCH_CUDA}"; exit 1 ;;
esac

# ------------------------------------------------------------
# 4. Clone or update repo
# ------------------------------------------------------------
echo ""
echo "[4/10] Downloading trainer repo..."
mkdir -p /workspace
if [ ! -d "${APP_DIR}/.git" ]; then
  git clone "${REPO_URL}" "${APP_DIR}"
else
  git -C "${APP_DIR}" pull
fi
cd "${APP_DIR}"

# ------------------------------------------------------------
# 5. Patch app for RunPod and selected precision
# ------------------------------------------------------------
echo ""
echo "[5/10] Patching app.py for RunPod and selected GPU precision..."
python3 - <<PYCODE
from pathlib import Path

precision = "${ANIMA_MIXED_PRECISION}"
p = Path("app.py")
t = p.read_text(encoding="utf-8")
t = t.replace('server_name="127.0.0.1"', 'server_name="0.0.0.0"')
t = t.replace('"base_model": "anima-preview",', '"base_model": "anima-preview3-base",')
t = t.replace('"mixed_precision": "bf16",', f'"mixed_precision": "{precision}",')
if precision == "fp16":
    t = t.replace('"save_precision": "bf16",', '"save_precision": "fp16",')
p.write_text(t, encoding="utf-8")
PYCODE

mkdir -p app_configs
cat > app_configs/accelerate_gpu.yaml <<YAML
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: 'NO'
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: '${ANIMA_MIXED_PRECISION}'
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
use_cpu: false
YAML

# ------------------------------------------------------------
# 6. Python virtual environment
# ------------------------------------------------------------
echo ""
echo "[6/10] Creating Python virtual environment..."
if [ ! -d ".venv" ]; then
  ${PYTHON_BIN} -m venv .venv
else
  echo ".venv already exists, reusing it."
fi
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel

# ------------------------------------------------------------
# 7. Install selected PyTorch build
# ------------------------------------------------------------
echo ""
echo "[7/10] Installing PyTorch..."
echo "Command: ${TORCH_INSTALL_CMD}"
pip uninstall -y torch torchvision torchaudio 2>/dev/null || true
eval "${TORCH_INSTALL_CMD}"
python - <<'PYCODE'
import torch
print("PyTorch:", torch.__version__)
print("CUDA:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
print("GPU:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE")
PYCODE

# ------------------------------------------------------------
# 8. Install sd-scripts and app requirements
# ------------------------------------------------------------
echo ""
echo "[8/10] Installing sd-scripts and requirements..."
if [ ! -d "sd-scripts/.git" ]; then
  git clone https://github.com/kohya-ss/sd-scripts.git sd-scripts
else
  git -C sd-scripts pull
fi

pushd sd-scripts >/dev/null
pip install -r requirements.txt
popd >/dev/null

pip install -r requirements.txt

echo "Re-checking PyTorch after requirements..."
python - <<'PYCODE'
import torch
print("PyTorch:", torch.__version__)
print("CUDA:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
print("GPU:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE")
PYCODE

# ------------------------------------------------------------
# 9. Download required models
# ------------------------------------------------------------
echo ""
echo "[9/10] Downloading Anima base model files..."
mkdir -p models/anima/dit
mkdir -p models/anima/text_encoder
mkdir -p models/anima/vae

QWEN_PATH="models/anima/text_encoder/qwen_3_06b_base.safetensors"
VAE_PATH="models/anima/vae/qwen_image_vae.safetensors"
PREVIEW3_PATH="models/anima/dit/anima-base-v1.0.safetensors"

if [ ! -f "${QWEN_PATH}" ]; then
  echo "Downloading Qwen3 text encoder, about 1.19 GB..."
  wget -c --show-progress -O "${QWEN_PATH}" \
    "https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/text_encoders/qwen_3_06b_base.safetensors"
else
  echo "Qwen3 text encoder already exists, skipping."
fi

if [ ! -f "${VAE_PATH}" ]; then
  echo "Downloading Qwen Image VAE, about 254 MB..."
  wget -c --show-progress -O "${VAE_PATH}" \
    "https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/vae/qwen_image_vae.safetensors"
else
  echo "VAE already exists, skipping."
fi

if [ ! -f "${PREVIEW3_PATH}" ]; then
  echo "Downloading Anima Preview 3 base model, about 4 GB..."
  wget -c --show-progress -O "${PREVIEW3_PATH}" \
    "https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/diffusion_models/anima-base-v1.0.safetensors"
else
  echo "Anima base model already exists, skipping."
fi

# ------------------------------------------------------------
# 10. Create launcher
# ------------------------------------------------------------
echo ""
echo "[10/10] Creating RunPod launcher..."
cat > run_anima_base_runpod.sh <<RUNNER
#!/usr/bin/env bash
set -e
cd "${APP_DIR}"
source .venv/bin/activate
export PYTHONIOENCODING=utf-8
echo "Starting Anima Preview 3 LoRA Trainer on port ${PORT}..."
echo "Selected PyTorch CUDA build: ${ANIMA_TORCH_CUDA}"
echo "Selected mixed precision: ${ANIMA_MIXED_PRECISION}"
echo "Use RunPod Connect -> HTTP Service ${PORT}, or open: https://\${RUNPOD_POD_ID}-${PORT}.proxy.runpod.net"
python app.py
RUNNER
chmod +x run_anima_base_runpod.sh

cat > INSTALL_INFO.txt <<INFO
Anima base LoRA Trainer adaptive install

GPU: ${ANIMA_GPU_NAME}
Driver: ${ANIMA_DRIVER_VERSION}
nvidia-smi CUDA shown: ${ANIMA_NVIDIA_SMI_CUDA}
Selected PyTorch CUDA build: ${ANIMA_TORCH_CUDA}
Selected mixed precision: ${ANIMA_MIXED_PRECISION}
Reason: ${ANIMA_GPU_REASON}

Run command:
cd ${APP_DIR}
./run_anima_base_runpod.sh

RunPod:
Expose HTTP port ${PORT}, then open Connect -> HTTP Service ${PORT}.
INFO

echo ""
echo "============================================================"
echo " INSTALL COMPLETE!"
echo "============================================================"
echo ""
echo "GPU: ${ANIMA_GPU_NAME}"
echo "Driver: ${ANIMA_DRIVER_VERSION}"
echo "PyTorch CUDA build: ${ANIMA_TORCH_CUDA}"
echo "Default mixed precision: ${ANIMA_MIXED_PRECISION}"
echo ""
echo "IMPORTANT: In RunPod, expose HTTP port ${PORT}."
echo "Then open Connect -> HTTP Service ${PORT}."
echo ""
echo "Starting the trainer now..."
echo ""
./run_anima_base_runpod.sh