import os, pickle, pandas as pd, streamlit as st
from huggingface_hub import hf_hub_download
import logging
logging.basicConfig(level=logging.INFO)

MODEL_LOCAL_PATH = "models/best_model.pkl"
MODEL_REPO = os.environ.get("MODEL_REPO", "username/model-name")
HF_TOKEN = os.environ.get("HF_TOKEN", None)
REPO_TYPE = os.environ.get("HF_REPO_TYPE", "model")

def ensure_model():
    if os.path.exists(MODEL_LOCAL_PATH):
        return MODEL_LOCAL_PATH
    os.makedirs("models", exist_ok=True)
    p = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pkl", repo_type=REPO_TYPE, token=HF_TOKEN)
    with open(p, "rb") as r, open(MODEL_LOCAL_PATH, "wb") as w:
        w.write(r.read())
    return MODEL_LOCAL_PATH

@st.cache_resource
def load_model():
    path = ensure_model()
    with open(path, "rb") as f:
        model = pickle.load(f)
    return model

st.title("Model Inference (Streamlit)")
st.write("Upload a CSV file or paste JSON/CSV rows to get predictions.")

uploaded = st.file_uploader("Upload CSV", type=["csv"])
if uploaded is not None:
    df = pd.read_csv(uploaded)
    st.write("Input preview:", df.head())
    if st.button("Predict"):
        model = load_model()
        preds = model.predict(df)
        df["prediction"] = preds
        st.write(df)

text_input = st.text_area("Paste CSV text or JSON list of dicts", height=150)
if st.button("Predict from text") and text_input.strip():
    try:
        df2 = pd.read_csv(pd.io.common.StringIO(text_input))
    except Exception:
        import json
        df2 = pd.DataFrame(json.loads(text_input))
    st.write("Parsed input:", df2.head())
    model = load_model()
    preds = model.predict(df2)
    df2["prediction"] = preds
    st.write(df2)