import os
from pathlib import Path
from urllib.request import urlretrieve
import streamlit as st

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:
    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}"
        
        st.info(f"Downloading model from {olmo_model_url} to {olmo_model}")
        try:
            APP_CACHE.mkdir(parents=True, exist_ok=True)
            urlretrieve(olmo_model_url, olmo_model)
            st.success("Download complete!")
        except Exception as e:
            st.error(f"Failed to download model: {str(e)}")
            raise
        return olmo_model
    else:
        st.info(f"Model already exists at {OLMO_MODEL}")
        return OLMO_MODEL

if __name__ == "__main__":
    download_olmo_model()