import requests
import os
from huggingface_hub import hf_hub_download
from utils import custom_drive_cache_dir, get_drive

HF_TOKEN = os.getenv("HF_TOKEN")

MANGA_LINE_EXTRACTION_MODEL = "https://github.com/ljsabc/MangaLineExtraction_PyTorch/releases/download/v1/erika.pth"
ANIME2SKETCH_MODEL = {"REPO_ID": "p1atdev/Anime2Sketch", "FILENAME": "netG.pth"}


def download_manga_line_extraction_model():
    if os.path.exists("./models/erika.pth"):
        return

    with requests.get(MANGA_LINE_EXTRACTION_MODEL, stream=True) as r:
        r.raise_for_status()
        with open("./models/erika.pth", "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)


def download_anime2sketch_model():
    if os.path.exists("./models/netG.pth"):
        return

    drive = get_drive("./models/netG.pth")
    with custom_drive_cache_dir(drive) as cache_dir:
        hf_hub_download(
            repo_id=ANIME2SKETCH_MODEL["REPO_ID"],
            filename=ANIME2SKETCH_MODEL["FILENAME"],
            local_dir="./models",
            use_auth_token=HF_TOKEN,
            local_dir_use_symlinks=False,
            cache_dir=cache_dir,
        )


def setup():
    if not os.path.exists("./models"):
        os.makedirs("./models")
    download_manga_line_extraction_model()
    download_anime2sketch_model()