"""Put this folder and optional OmniVAE sources on the import path."""

import os
import sys
from pathlib import Path


def prepare() -> Path:
    repo = Path(__file__).resolve().parent
    os.chdir(repo)
    paths = [str(repo)]

    omni = os.environ.get("OMNIVAE_CODE_ROOT", "").strip()
    if not omni:
        fallback = repo.parent.parent / "model" / "OmniVAE_code"
        if fallback.is_dir():
            omni = str(fallback)
    if omni:
        paths.append(omni)
        os.environ.setdefault("OMNIVAE_CODE_ROOT", omni)

    if not os.environ.get("OMNIVAE_AUDIO_CKPT_DIR", "").strip():
        audio = (
            repo.parent.parent
            / "model"
            / "OmniVAE"
            / "models"
            / "vae"
            / "audio_only"
            / "recon_distill_avclip_ft_decoder"
        )
        if audio.is_dir():
            os.environ.setdefault("OMNIVAE_AUDIO_CKPT_DIR", str(audio))

    for name in (
        "DAVINCI_DISABLE_MAGI_COMPILE",
        "MAGI_DISABLE_COMPILE",
        "TORCHDYNAMO_DISABLE",
    ):
        os.environ.setdefault(name, "1")

    sys.path[:0] = paths
    return repo