import os
import subprocess
import modal

volume = modal.Volume.from_name("Output", create_if_missing = True)

image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git", "libgl1", "libglib2.0-0")
    .add_local_file(local_path = "config.json",    remote_path = "/root/config.json")
    .add_local_dir(local_path  = "thirddiffusion", remote_path = "/root/thirddiffusion")
)

app = modal.App(
    name    = "OneTrainer",
    image   = image,
    volumes = {"/workspace": volume}
)

@app.function(gpu = "A100", timeout = 14400)
def train_lora():

    WORKDIR = "/workspace"

    REPODIR = os.path.join(WORKDIR, "OneTrainer")

    CONFIG_PATH = "/root/config.json"

    DATA_PATH = "/root/thirddiffusion"

    try:

        if not os.path.exists(REPODIR):

            subprocess.run(
                ["git", "clone", "https://github.com/Nerogar/OneTrainer.git", REPODIR],
                check  = True,
                stderr = subprocess.PIPE
            )

            INSTALL_PATH = os.path.join(REPODIR, "install.sh")

            subprocess.run(
                ["bash", INSTALL_PATH],
                cwd   = REPODIR,
                check = True,
                env   = dict(os.environ, DEBIAN_FRONTEND = "noninteractive")
            )

        else:

            if os.path.exists(os.path.join(REPODIR, "out")):

                subprocess.run(["rm", "-rf", os.path.join(REPODIR, "out")], check = True)

            if os.path.exists(os.path.join(REPODIR, "tmp")):

                subprocess.run(["rm", "-rf", os.path.join(REPODIR, "tmp")], check = True)

        PYTHON_PATH = os.path.join(REPODIR, "venv", "bin", "python")

        #subprocess.run(
        #    [PYTHON_PATH, "-m", "pip", "install", "-r", os.path.join(REPODIR, "requirements.txt")],
        #    check = True
        #)

        subprocess.run(
            [
                PYTHON_PATH,
                os.path.join(REPODIR, "scripts", "train.py"),
                "--config-path",
                CONFIG_PATH
            ],
            cwd   = REPODIR,
            check = True,
            env   = dict(os.environ, CUDA_VISIBLE_DEVICES = "0")
        )

        volume.commit()

    except subprocess.CalledProcessError as e:

        raise RuntimeError(f"Command failed: {e.stderr.decode()}") from e

@app.local_entrypoint()
def main():

    try:

        result = train_lora.remote()

        print(f"Success: {result}")

    except modal.exception.ExecutionError as e:

        print(f"Failed: {e}")

        print("Check logs with: modal logs " + app.app_id)

if __name__ == "__main__":

    main()