import os
from pathlib import Path
from urllib.request import urlretrieve

HOME_DIR = Path.home()
APP_CACHE = HOME_DIR / "olmo_demo"

OLMO_MODEL_FILE = os.environ.get("OLMO_MODEL_FILE", "OLMo-7B-Instruct-Q4_K_M.gguf")
OLMO_MODEL = APP_CACHE / OLMO_MODEL_FILE

def download_olmo_model(model_file: str | None = None, force=False) -> Path:
    """Download the OLMO model from the Hugging Face model hub.

    Parameters
    ----------
    model_file : str | None, optional
        The name of the model file to download, by default None
    force : bool, optional
        Whether to force the download even if the file already exists, by default False

    Returns
    -------
    pathlib.Path
        The path to the downloaded model file
    """
    if not OLMO_MODEL.exists() or force:
        if model_file is None:
            model_file = OLMO_MODEL_FILE
            olmo_model = OLMO_MODEL
        else:
            olmo_model = APP_CACHE / model_file
        olmo_model_url = f"https://huggingface.co/ssec-uw/OLMo-7B-Instruct-GGUF/resolve/main/{model_file}"
        print(f"Downloading model from {olmo_model_url} to {olmo_model}")
        urlretrieve(olmo_model_url, olmo_model)
        print("Download complete!")
        return olmo_model
    else:
        print(f"Model already exists at {OLMO_MODEL}")
        return OLMO_MODEL

if __name__ == "__main__":
    APP_CACHE.mkdir(parents=True, exist_ok=True)
    download_olmo_model()