#!/usr/bin/env bash
# install_onetrainer_intel_xpu.sh
# Reinstall OneTrainer for Intel XPU (CachyOS / no containers).
#
# Based on the OneTrainer section of updateai (PR 1413 + XPU patches).
#
# Requires: Bash 5+, git, curl, uv (installed automatically),
#           Intel oneAPI Base Toolkit at /opt/intel/oneapi/setvars.sh
#
# Usage:
#   chmod +x install_onetrainer_intel_xpu.sh
#   ./install_onetrainer_intel_xpu.sh
#
# Launch after install:
#   source ${HOME}/ai/onetrainer/venv/bin/activate
#   python ${HOME}/ai/onetrainer/scripts/train_ui_qt.py
#
# Note: this script recreates the venv. If you use make_masks.py in the
# same environment, reinstall rembg afterwards:
#   source ${HOME}/ai/onetrainer/venv/bin/activate
#   pip install rembg pillow numpy
set -euo pipefail

readonly BASEDIR="${HOME}/ai"
readonly OT="${BASEDIR}/onetrainer"
readonly PYTHON_VER="3.13"
readonly UV_INSTALL_URL="https://astral.sh/uv/install.sh"
readonly PYTORCH_XPU_STABLE_INDEX="https://download.pytorch.org/whl/xpu"
readonly ONETRAINER_REPO="https://github.com/Nerogar/OneTrainer.git"
readonly ONETRAINER_PR="1413"
readonly ONETRAINER_DIFFUSERS_REF="1ffa423"
readonly ONETRAINER_MGDS_REF="3a6994a"
readonly ONETRAINER_MUON_REF="f90a42b"
readonly ONEAPI_SETVARS="/opt/intel/oneapi/setvars.sh"

die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; }
log() { printf '\n=== %s ===\n' "$*"; }

ensure_uv() {
    export PATH="${HOME}/.local/bin:${PATH}"
    if command -v uv >/dev/null 2>&1; then
        return 0
    fi
    log "Installing uv"
    curl -LsSf "${UV_INSTALL_URL}" | sh
    export PATH="${HOME}/.local/bin:${PATH}"
}

ot_uv() { UV_NO_PROJECT=1 uv "$@"; }

sync_src() {
    mkdir -p "$(dirname "${OT}")"
    if [[ -d "${OT}/.git" ]]; then
        log "Updating OneTrainer source"
        git -C "${OT}" fetch --prune origin
        git -C "${OT}" checkout --force master
        git -C "${OT}" reset --hard origin/master
        git -C "${OT}" clean -fd --exclude=venv --exclude=workspace --exclude=workspace-cache
    else
        log "Cloning OneTrainer → ${OT}"
        rm -rf "${OT}"
        git clone --quiet "${ONETRAINER_REPO}" "${OT}"
    fi

    log "Merging PR ${ONETRAINER_PR}"
    git -C "${OT}" fetch origin "pull/${ONETRAINER_PR}/head:pr-${ONETRAINER_PR}"
    git -C "${OT}" merge "pr-${ONETRAINER_PR}" --no-edit || true
    if [[ -f "${OT}/.git/MERGE_HEAD" ]]; then
        if git -C "${OT}" diff --name-only --diff-filter=U | grep -qx 'scripts/train_ui.py'; then
            git -C "${OT}" checkout --ours scripts/train_ui.py
            git -C "${OT}" add scripts/train_ui.py
        fi
        if git -C "${OT}" diff --name-only --diff-filter=U | grep -q .; then
            die "Unresolved OneTrainer merge conflicts remain"
        fi
        git -C "${OT}" add -u
        git -C "${OT}" commit --no-edit
    fi
}

patch_xpu() {
    log "Applying XPU patches"
    python - <<PY
from pathlib import Path
import re

root = Path("${OT}")

p = root / "scripts/util/import_util.py"
text = p.read_text()
if "ipex_init" not in text:
    p.write_text(text.rstrip() + """

    import torch

    if hasattr(torch, "xpu") and torch.xpu.is_available():
        from ipex_to_cuda import ipex_init

        ipex_init()
        print("CUDA -> XPU hijacking active")
""")
    print("import_util.py: ipex hook added")
else:
    print("import_util.py: ipex hook already present")

tu = root / "modules/util/torch_util.py"
kept, removed = [], 0
for line in tu.read_text().splitlines(keepends=True):
    if 'print ("true")' in line or 'print("true")' in line:
        removed += 1
        continue
    kept.append(line)
if removed:
    tu.write_text("".join(kept))
    print(f"torch_util.py: removed {removed} debug print(s)")
else:
    print("torch_util.py: no debug print")

zluda = root / "modules/zluda/ZLUDA.py"
ztext = zluda.read_text()
new_fn = '''def is_zluda(device):
    try:
        if isinstance(device, str):
            if device.lower() == "cpu":
                return False
            if device.lower().startswith("cuda"):
                parts = device.split(":")
                device = int(parts[1]) if len(parts) > 1 else 0
        return str(torch.cuda.get_device_name(device)).endswith("[ZLUDA]")
    except Exception:
        return False
'''
ztext2, n = re.subn(
    r"def is_zluda\([\s\S]*?(?=\ndef |\Z)",
    new_fn + "\n",
    ztext,
    count=1,
)
if n != 1:
    raise SystemExit(f"ZLUDA is_zluda replace count={n}")
zluda.write_text(ztext2)
print("ZLUDA.py: is_zluda patched")
PY
    python -m py_compile \
        "${OT}/scripts/util/import_util.py" \
        "${OT}/modules/util/torch_util.py" \
        "${OT}/modules/zluda/ZLUDA.py"
}

install_python() {
    [[ -f "${ONEAPI_SETVARS}" ]] || die "oneAPI not found at ${ONEAPI_SETVARS}"

    log "Creating Python ${PYTHON_VER} venv"
    rm -rf "${OT}/venv"
    ot_uv venv "${OT}/venv" --python "${PYTHON_VER}"
    # shellcheck source=/dev/null
    source "${OT}/venv/bin/activate"

    log "Installing torch XPU + ipex_to_cuda"
    ot_uv pip install torch torchvision torchaudio --index-url "${PYTORCH_XPU_STABLE_INDEX}"
    ot_uv pip install "git+https://github.com/Disty0/ipex_to_cuda.git"
    ot_uv pip install \
        "git+https://github.com/huggingface/diffusers.git@${ONETRAINER_DIFFUSERS_REF}" \
        "git+https://github.com/Nerogar/mgds.git@${ONETRAINER_MGDS_REF}" \
        "git+https://github.com/KellerJordan/Muon.git@${ONETRAINER_MUON_REF}"
    grep -vE '^\-e |^numpy==' "${OT}/requirements-global.txt" | ot_uv pip install -r -
    ot_uv pip install onnxruntime

    python - <<'PY'
import torch
print("torch", torch.__version__)
print("xpu", torch.xpu.is_available(), "count", torch.xpu.device_count())
if not torch.xpu.is_available():
    raise SystemExit("XPU not available in OneTrainer venv")
print("device", torch.xpu.get_device_name(0))
PY
    deactivate
}

main() {
    log "Installing OneTrainer XPU → ${OT}"
    mkdir -p "${BASEDIR}"
    ensure_uv
    sync_src
    patch_xpu
    install_python

    printf '\nReady: %s\n' "${OT}"
    printf 'Launch:\n  source %s/venv/bin/activate\n  python %s/scripts/train_ui_qt.py\n' "${OT}" "${OT}"
    printf 'Kept: %s/workspace  %s/workspace-cache\n' "${OT}" "${OT}"
}

main "$@"