from __future__ import annotations

import modal

from config import cnf
from embedding_train.image import (
    data_volume,
    hf_cache_volume,
    training_image,
)

app = modal.App("prepare-garments2look")


@app.function(
    image=training_image,
    cpu=4,
    memory=16384,
    timeout=4 * 60 * 60,
    volumes={
        cnf.train_data_mnt: data_volume,
        cnf.train_hf_cache_mnt: hf_cache_volume,
    },
)
def download():
    import shutil
    import tarfile
    from pathlib import Path

    from huggingface_hub import HfApi, hf_hub_download

    repo_id = "ArtmeScienceLab/Garments2Look"
    destination = Path(cnf.train_data_mnt) / "Garments2Look"
    marker = destination / ".complete"

    if marker.exists():
        print(f"Dataset already prepared at {destination}")
        return

    destination.mkdir(parents=True, exist_ok=True)

    api = HfApi()
    repo_files = api.list_repo_files(
        repo_id=repo_id,
        repo_type="dataset",
    )

    wanted_names = {
        "annotations.tar.gz",
        "images.tar.gz",
    }

    archive_names = [
        filename
        for filename in repo_files
        if filename.startswith("polyvore/")
        and Path(filename).name in wanted_names
    ]

    if len(archive_names) != 2:
        raise RuntimeError(
            "Could not locate the expected Polyvore archives.\n"
            f"Found: {archive_names}"
        )

    for filename in archive_names:
        print(f"Downloading {filename}")

        archive_path = hf_hub_download(
            repo_id=repo_id,
            repo_type="dataset",
            filename=filename,
        )

        print(f"Extracting {filename}")

        with tarfile.open(archive_path, "r:gz") as archive:
            archive.extractall(destination)

    marker.touch()
    data_volume.commit()

    print(f"Dataset prepared at {destination}")


@app.local_entrypoint()
def main():
    download.remote()