import os
import base64
import requests
import gradio as gr
import spaces
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification, pipeline

# --- Roboflow face-detection key, read from a Space secret (not hardcoded) ---
ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")

# --- Load all three classifiers once at startup ---
# These run locally in this container via `transformers`, instead of calling
# Hugging Face's shared serverless Inference API (which no longer serves the
# gender/race models — only age currently has an attached provider).

def load_vit_pipeline(model_name):
    """
    Load explicitly as ViT rather than relying on AutoImageProcessor/AutoModel.
    These models were fine-tuned ~2 years ago using the standard HF ViT
    image-classification recipe; newer `transformers` versions sometimes
    fail to auto-detect older model configs, so we bypass Auto* entirely.
    """
    processor = ViTImageProcessor.from_pretrained(model_name)
    model = ViTForImageClassification.from_pretrained(model_name)
    return pipeline("image-classification", model=model, image_processor=processor)

print("Loading race/age/gender models — this can take a bit on first boot...")
race_pipe = load_vit_pipeline("crangana/trained-race")
age_pipe = load_vit_pipeline("crangana/trained-age")
gender_pipe = load_vit_pipeline("crangana/trained-gender")
print("Models loaded.")


def crop_face(image_path):
    """Detect and crop the face using Roboflow's hosted face detector."""
    if not ROBOFLOW_API_KEY:
        raise ValueError(
            "ROBOFLOW_API_KEY is not set. Add it under this Space's "
            "Settings → Repository secrets."
        )

    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")

    url = "https://detect.roboflow.com/face-det-npvna/1"
    params = {"api_key": ROBOFLOW_API_KEY}
    headers = {"Content-Type": "application/x-www-form-urlencoded"}

    response = requests.post(url, params=params, data=image_base64, headers=headers)

    try:
        crop_data = response.json()
    except ValueError:
        raise ValueError(
            f"Roboflow returned a non-JSON response (status {response.status_code}): "
            f"{response.text[:300]}"
        )

    print("Roboflow response:", crop_data)  # visible in the Space's Logs tab

    if "error" in crop_data:
        raise ValueError(f"Roboflow error: {crop_data['error']}")

    if response.status_code != 200:
        raise ValueError(f"Roboflow returned status {response.status_code}: {crop_data}")

    if not crop_data.get("predictions"):
        raise ValueError(
            "Roboflow ran successfully but found no face in this image. "
            "Try a clearer, front-facing photo."
        )

    pred = crop_data["predictions"][0]
    width, height = pred["width"], pred["height"]
    x = pred["x"] - width / 2
    y = pred["y"] - height / 2

    image = Image.open(image_path)
    cropped = image.crop((x, y, x + width, y + height))

    # Force an exact 224x224 square before it reaches the models — the
    # Roboflow bounding box isn't guaranteed to be square, and this
    # particular model's built-in preprocessing doesn't fully correct
    # for that on its own (was producing e.g. 227x224 instead of 224x224,
    # which the model's fixed patch grid can't accept).
    cropped = cropped.convert("RGB").resize((224, 224), Image.LANCZOS)

    cropped_path = "/tmp/cropped_face.jpg"
    cropped.save(cropped_path)
    return cropped_path


def get_top(results):
    """Pick the highest-confidence label from a list of {label, score} dicts."""
    best = max(results, key=lambda r: r["score"])
    return best["label"], best["score"]


@spaces.GPU
def classify(image):
    if image is None:
        return "No image provided.", "", ""

    tmp_path = "/tmp/input_face.jpg"
    image.convert("RGB").save(tmp_path)

    try:
        cropped_path = crop_face(tmp_path)
    except Exception as e:
        return f"Face detection failed: {e}", "—", "—"

    try:
        age_label, age_score = get_top(age_pipe(cropped_path))
        gender_label, gender_score = get_top(gender_pipe(cropped_path))
        race_label, race_score = get_top(race_pipe(cropped_path))
    except Exception as e:
        return f"Classification failed: {e}", "—", "—"

    return (
        f"{age_label} ({age_score:.0%})",
        f"{gender_label} ({gender_score:.0%})",
        f"{race_label} ({race_score:.0%})",
    )


demo = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="Upload a face photo"),
    outputs=[
        gr.Textbox(label="Age"),
        gr.Textbox(label="Gender"),
        gr.Textbox(label="Ethnicity"),
    ],
    title="Face Attribute Classifier",
    description="Predicts age, gender, and ethnicity from a face photo.",
    api_name="classify",  # names the programmatic API endpoint /gradio_api/call/classify
)

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