import os
import importlib.util
from huggingface_hub import hf_hub_download
import gradio as gr

# --- CONFIG ---
PRIVATE_DATASET_ID = "abdulrafay9/containeralign-private"
TOKEN = os.environ.get("HF_TOKEN")

if not TOKEN:
    raise RuntimeError("HF_TOKEN is not set. Add it in Settings → Variables and secrets → Secrets.")

# --- DOWNLOAD FILES ---
core_path = hf_hub_download(
    repo_id=PRIVATE_DATASET_ID,
    repo_type="dataset",
    filename="app_core.py",
    token=TOKEN,
)
weights_path = hf_hub_download(
    repo_id=PRIVATE_DATASET_ID,
    repo_type="dataset",
    filename="model.pth",
    token=TOKEN,
)

# --- LOAD MODULE ---
spec = importlib.util.spec_from_file_location("app_core", core_path)
app_core = importlib.util.module_from_spec(spec)
spec.loader.exec_module(app_core)

# --- LOAD MODEL ---
model = app_core.load_model(weights_path)

# --- DEFINE PREDICTION FUNCTION ---
def predict(image):
    return app_core.predict_alignment(model, image)

# --- GRADIO INTERFACE ---
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Textbox(label="Prediction Result"),
    title="Container Alignment Detection",
    description="Upload an image to check whether containers are Aligned or Not Aligned."
)

# --- RUN APP ---
if __name__ == "__main__":
    demo.launch()