import os
from typing import Dict
from PIL import Image

import torch
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoConfig

# ----- CONFIG -----
MODEL_ID = os.getenv("MODEL_ID", "kritaphatson/thai_snake_image_classifier").strip()
TOPK = int(os.getenv("TOPK", "3"))
torch.set_num_threads(max(1, int(os.getenv("NUM_THREADS", "1"))))

# ----- LOAD -----
# Disable fast processor to avoid the torchvision warning on CPU Spaces
processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_fast=False)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
model.eval().to("cpu")

cfg = AutoConfig.from_pretrained(MODEL_ID)
id2label: Dict[int, str] = {int(k): v for k, v in getattr(cfg, "id2label", {}).items()} or {
    i: f"class_{i}" for i in range(model.config.num_labels)
}

# ----- PREDICT -----
@torch.inference_mode()
def predict(img: Image.Image) -> Dict[str, float]:
    if img.mode != "RGB":
        img = img.convert("RGB")
    inputs = processor(images=img, return_tensors="pt")
    probs = model(**inputs).logits.softmax(dim=-1)[0]
    values, indices = torch.topk(probs, k=min(TOPK, probs.numel()))
    return {id2label[int(i)]: float(v) for v, i in zip(values, indices)}

# ----- UI -----
DESCRIPTION = """
Upload a snake image. The model returns the top species predictions.
Model: **kritaphatson/thai_snake_image_classifier**.
"""

with gr.Blocks(title="Thai Snake Classifier") as demo:
    gr.Markdown("# Thai Snake Classifier")
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload image", sources=["upload", "clipboard"])
        out = gr.Label(label=f"Top {TOPK} predictions")
    gr.Button("Classify").click(fn=predict, inputs=inp, outputs=out)

if __name__ == "__main__":
    # Gradio 4: queue() takes no concurrency args
    demo.queue(max_size=8).launch()