import os
from functools import lru_cache

import gradio as gr
import numpy as np
from PIL import Image


MODEL_INFO_PATH = os.getenv("MODEL_INFO_PATH", "model_info.json")
ENV_MODEL_PATH = os.getenv("MODEL_PATH")

# INPUT_SCALE options:
# - "0_1":   image/255.0  (common for custom CNNs)
# - "minus1_1": image/127.5 - 1.0 (common for MobileNet-style models)
# - "0_255": keep 0..255 (less common, but sometimes used)
ENV_INPUT_SCALE = os.getenv("INPUT_SCALE")

ENV_CLASS_NAMES = os.getenv("CLASS_NAMES")
ENV_POSITIVE_CLASS = os.getenv("POSITIVE_CLASS")


def _patch_gradio_client_schema_bool() -> None:
    try:
        import gradio_client.utils as gcu

        original = getattr(gcu, "_json_schema_to_python_type", None)
        if not callable(original):
            return

        def patched(schema, defs=None):
            if isinstance(schema, bool):
                return "Any"
            return original(schema, defs)

        gcu._json_schema_to_python_type = patched
    except Exception:
        return


_patch_gradio_client_schema_bool()


def _parse_class_names(text: str) -> list[str]:
    return [name.strip() for name in (text or "").split(",") if name.strip()]


def _load_model_info() -> dict:
    if not MODEL_INFO_PATH or not os.path.exists(MODEL_INFO_PATH):
        return {}
    try:
        import json

        with open(MODEL_INFO_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


_MODEL_INFO = _load_model_info()


def _coalesce(*values: str | None) -> str | None:
    for value in values:
        if value is None:
            continue
        value = value.strip()
        if value:
            return value
    return None


def _guess_model_path() -> str:
    if ENV_MODEL_PATH:
        return ENV_MODEL_PATH

    for key in ("model_path", "model_file", "model", "artifact"):
        value = _MODEL_INFO.get(key)
        if isinstance(value, str) and value.strip() and os.path.exists(value.strip()):
            return value.strip()

    if os.path.exists("model.h5"):
        return "model.h5"

    h5_files = [p for p in os.listdir(".") if p.lower().endswith(".h5")]
    if len(h5_files) == 1:
        return h5_files[0]

    return "model.h5"


def _guess_class_names_default() -> str:
    if ENV_CLASS_NAMES is not None:
        return ENV_CLASS_NAMES.strip()

    for key in ("class_names", "classes", "labels"):
        value = _MODEL_INFO.get(key)
        if isinstance(value, list) and all(isinstance(x, str) for x in value):
            return ",".join([x.strip() for x in value if x.strip()])
        if isinstance(value, str) and value.strip():
            return value.strip()

    return "cat,dog"


def _guess_input_scale_default() -> str:
    if ENV_INPUT_SCALE is not None:
        return ENV_INPUT_SCALE.strip()

    for key in ("input_scale", "scale", "preprocess", "normalization"):
        value = _MODEL_INFO.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()

    model_path = str(MODEL_PATH).lower()
    if "mobilenetv2" in model_path or "mobilenet_v2" in model_path:
        return "minus1_1"

    return "0_1"


def _guess_positive_class_default(class_names_text: str) -> str:
    if ENV_POSITIVE_CLASS is not None:
        return ENV_POSITIVE_CLASS.strip()

    value = _MODEL_INFO.get("positive_class")
    if isinstance(value, str) and value.strip():
        return value.strip()

    names = _parse_class_names(class_names_text)
    if len(names) >= 2:
        return names[1]
    return "dog"


MODEL_PATH = _guess_model_path()
DEFAULT_CLASS_NAMES = _guess_class_names_default()
DEFAULT_INPUT_SCALE = _guess_input_scale_default()
DEFAULT_POSITIVE_CLASS = _guess_positive_class_default(DEFAULT_CLASS_NAMES)


def _normalize_image(x: np.ndarray, input_scale: str) -> np.ndarray:
    x = x.astype(np.float32)
    if input_scale == "0_1":
        return x / 255.0
    if input_scale == "minus1_1":
        return (x / 127.5) - 1.0
    if input_scale == "0_255":
        return x
    raise ValueError(f"Unknown INPUT_SCALE={input_scale!r} (use 0_1, minus1_1, or 0_255)")


@lru_cache(maxsize=1)
def _load_model():
    try:
        import tensorflow as tf
    except Exception as e:  # pragma: no cover
        raise RuntimeError(
            "TensorFlow import failed. Ensure `tensorflow-cpu` is installed via requirements.txt."
        ) from e

    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(
            f"Model file not found at {MODEL_PATH!r}. Put your .h5 in the Space repo root as `model.h5`, "
            "or set the MODEL_PATH environment variable in the Space settings."
        )

    return tf.keras.models.load_model(MODEL_PATH, compile=False)


def _get_target_size(model) -> tuple[int, int]:
    shape = getattr(model, "input_shape", None)
    if not shape or len(shape) < 3:
        return (224, 224)

    h, w = shape[1], shape[2]
    if isinstance(h, int) and isinstance(w, int):
        return (w, h)
    return (224, 224)


def _to_probs(raw: np.ndarray) -> np.ndarray:
    raw = np.asarray(raw).reshape(-1).astype(np.float32)
    if raw.size == 1:
        v = float(raw[0])
        if 0.0 <= v <= 1.0:
            p = v
        else:
            p = float(1.0 / (1.0 + np.exp(-v)))
        return np.array([1.0 - p, p], dtype=np.float32)

    if raw.size == 2:
        s = float(raw.sum())
        if np.all(raw >= 0.0) and np.all(raw <= 1.0) and abs(s - 1.0) < 1e-3:
            return raw.astype(np.float32)
        shifted = raw - np.max(raw)
        exp = np.exp(shifted)
        probs = exp / np.sum(exp)
        return probs.astype(np.float32)

    s = float(raw.sum())
    if np.all(raw >= 0.0) and np.all(raw <= 1.0) and abs(s - 1.0) < 1e-3:
        return raw.astype(np.float32)
    shifted = raw - np.max(raw)
    exp = np.exp(shifted)
    probs = exp / np.sum(exp)
    return probs.astype(np.float32)


def predict(image: Image.Image, class_names_text: str, input_scale: str, positive_class: str) -> dict:
    try:
        model = _load_model()
        target_w, target_h = _get_target_size(model)

        class_names = _parse_class_names(class_names_text)
        if not class_names:
            class_names = ["cat", "dog"]

        image = image.convert("RGB").resize((target_w, target_h))
        arr = np.array(image)
        arr = _normalize_image(arr, input_scale=input_scale)
        arr = np.expand_dims(arr, axis=0)

        raw = model.predict(arr, verbose=0)
        probs = _to_probs(raw)

        if probs.size == 2 and len(class_names) == 2:
            if np.asarray(raw).reshape(-1).size == 1 and positive_class == class_names[0]:
                probs = probs[::-1]
            return {class_names[0]: float(probs[0]), class_names[1]: float(probs[1])}

        names = class_names if len(class_names) == probs.size else [f"class_{i}" for i in range(probs.size)]
        return {name: float(p) for name, p in zip(names, probs)}
    except Exception as e:
        raise gr.Error(str(e))


with gr.Blocks() as demo:
    gr.Markdown("# Cats vs Dogs Classifier")
    gr.Markdown(
        "Upload an image and the model will predict whether it looks like a cat or a dog.\n\n"
        f"- `MODEL_PATH`: `{MODEL_PATH}`\n"
        f"- `MODEL_INFO_PATH`: `{MODEL_INFO_PATH}`\n"
        f"- Default `CLASS_NAMES`: `{DEFAULT_CLASS_NAMES}`\n"
        f"- Default `INPUT_SCALE`: `{DEFAULT_INPUT_SCALE}`\n"
        f"- Default `POSITIVE_CLASS`: `{DEFAULT_POSITIVE_CLASS}`"
    )

    with gr.Accordion("Advanced settings", open=False):
        class_names_text = gr.Textbox(
            label="Class names (comma-separated)",
            value=DEFAULT_CLASS_NAMES,
            placeholder="cat,dog",
        )
        input_scale = gr.Dropdown(
            label="Input scaling",
            value=DEFAULT_INPUT_SCALE,
            choices=["0_1", "minus1_1", "0_255"],
        )
        positive_class = gr.Textbox(
            label="Positive class (only for 1-output sigmoid models)",
            value=DEFAULT_POSITIVE_CLASS,
            placeholder="dog",
        )

    with gr.Row():
        inp = gr.Image(type="pil", label="Input image")
        out = gr.Label(num_top_classes=2, label="Prediction")

    btn = gr.Button("Predict")
    btn.click(fn=predict, inputs=[inp, class_names_text, input_scale, positive_class], outputs=out)

    gr.Examples(
        examples=[],
        inputs=inp,
        label="Examples (optional: add sample images to the repo and list them here)",
    )

if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    demo.launch(server_name="0.0.0.0", server_port=port)