import io, os, tempfile, pathlib, zipfile, shutil
from PIL import Image, ImageOps
import pandas as pd  # needed to hold predictor output, but we won't return any DataFrame
from huggingface_hub import hf_hub_download
import autogluon.multimodal as ag
import gradio as gr

MODEL_REPO_ID = "its-zion-18/sign-image-autogluon-predictor"
ZIP_FILENAME  = "autogluon_image_predictor_dir.zip"
CACHE_DIR   = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"

CLASS_LABELS = {0: "Other sign", 1: "Stop sign"}
DEFAULT_SIZE = 224

def _prepare_predictor_dir() -> str:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    local_zip = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
        local_dir=str(CACHE_DIR),
        local_dir_use_symlinks=False,
    )
    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(str(EXTRACT_DIR))
    contents = list(EXTRACT_DIR.iterdir())
    return str(contents[0] if len(contents) == 1 and contents[0].is_dir() else EXTRACT_DIR)

PREDICTOR = ag.MultiModalPredictor.load(_prepare_predictor_dir())

def preprocess_for_model(img: Image.Image, size: int = DEFAULT_SIZE) -> Image.Image:
    return ImageOps.fit(img, (size, size), method=Image.BICUBIC, centering=(0.5, 0.5))

# -------------------------
# Prediction (return 4 values to match your outputs)
# -------------------------
def predict(img: Image.Image):
    if img is None:
        return None, None, {"Error": 1.0}, "**Error:** No image provided."
    pre_img = preprocess_for_model(img, DEFAULT_SIZE)

    # temp save for AutoGluon
    tmpdir = pathlib.Path(tempfile.mkdtemp())
    img_path = tmpdir / "input.png"
    pre_img.save(img_path)

    df = pd.DataFrame({"image": [str(img_path)]})
    proba_df = PREDICTOR.predict_proba(df)
    row = proba_df.iloc[0]

    # Handle int or str column names
    p0 = float(row.get(0, row.get("0", 0.0)))
    p1 = float(row.get(1, row.get("1", 0.0)))

    pretty = {CLASS_LABELS[0]: p0, CLASS_LABELS[1]: p1}
    status = ""  # empty if OK

    # EXACTLY 4 returns to match your outputs list
    return img, pretty, status

EXAMPLES = [
    "https://drive.google.com/uc?id=1OwjdQdiclDOCuUirVsCX-_J1q1Y_G-Xs",
    "https://drive.google.com/uc?id=1g9XwRJGcqYwJ_mFtHj25ZE1w2VZPkBiR",
    "https://drive.google.com/uc?id=1MH5aIvHWKdamwuzdQPWlNwaoRE3c2bVb"
]

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Input image (PNG/JPG)", sources=["upload", "webcam"]),
    outputs=[
        gr.Image(label="Preprocessed"),
        gr.Label(num_top_classes=2, label="Class probabilities"),
        gr.Markdown(label="Status"),
    ],
    title="🪧 Sign Identification",
    description="Upload a PNG/JPG. Predicts Stop vs Other sign.",
    allow_flagging="never",
    examples=EXAMPLES,
)

if __name__ == "__main__":
    demo.launch()