import io
import os
import shutil
import tempfile
from typing import Dict, List, Tuple

import numpy as np
import gradio as gr
from PIL import Image
import tensorflow as tf
from huggingface_hub import hf_hub_download

# =========================
# Config (via Variables)
# =========================
# Suporta três formas de apontar o modelo:
# 1) MODEL_URL (http/https) -> baixa direto por URL
# 2) MODEL_REPO + MODEL_FILE (+ MODEL_REPO_TYPE: model|space)
# 3) Caminho local (MODEL_FILE existente na raiz do Space)
MODEL_URL = os.environ.get("MODEL_URL", "").strip()
MODEL_REPO = os.environ.get("MODEL_REPO", "").strip()  # ex: "vcollos/raspagemTF" ou "spaces/vcollos/raspagem_supra"
MODEL_REPO_TYPE = os.environ.get("MODEL_REPO_TYPE", "model").strip()  # "model" ou "space"
MODEL_FILE = os.environ.get("MODEL_FILE", "raspagem_model_v1.pb").strip()
LABELS_FILE = os.environ.get("LABELS_FILE", "labels.txt").strip()
IMG_SIZE = int(os.environ.get("IMG_SIZE", "224"))
TOPK = int(os.environ.get("TOPK", "5"))

# =========================
# Baixa/resolve o SavedModel (.pb) e carrega via lazy init
# =========================

def _download_from_url(url: str) -> str:
    import requests
    resp = requests.get(url, timeout=60)
    resp.raise_for_status()
    tmp_dir = tempfile.mkdtemp(prefix="raspagem_dl_")
    local = os.path.join(tmp_dir, os.path.basename(url) or "saved_model.pb")
    with open(local, "wb") as f:
        f.write(resp.content)
    return local


def _download_model() -> str:
    # Prioridade: URL explícita -> HF repo -> arquivo local
    if MODEL_URL:
        return _download_from_url(MODEL_URL)

    if MODEL_REPO:
        try:
            return hf_hub_download(
                repo_id=MODEL_REPO,
                filename=MODEL_FILE,
                repo_type=MODEL_REPO_TYPE if MODEL_REPO_TYPE in {"model", "space"} else "model",
            )
        except Exception as e:
            print(f"[download] HF hub falhou: {e}")

    if os.path.exists(MODEL_FILE):
        return MODEL_FILE

    raise FileNotFoundError(
        "Modelo não encontrado. Defina MODEL_URL OU (MODEL_REPO, MODEL_REPO_TYPE, MODEL_FILE) OU deixe o arquivo na raiz do Space."
    )


def _prepare_saved_model_dir(pb_path: str) -> str:
    # SavedModel mínimo: diretório contendo 'saved_model.pb'
    tmp_dir = tempfile.mkdtemp(prefix="raspagem_savedmodel_")
    shutil.copy(pb_path, os.path.join(tmp_dir, "saved_model.pb"))
    return tmp_dir


# Lazy state
_SERVING_FN = None
_LABELS: List[str] = []
_LAST_INIT_ERROR: str | None = None


def _maybe_labels() -> List[str]:
    # Tenta arquivo labels.txt no HF repo/local
    try:
        if LABELS_FILE:
            if MODEL_REPO:
                p = hf_hub_download(
                    repo_id=MODEL_REPO,
                    filename=LABELS_FILE,
                    repo_type=MODEL_REPO_TYPE if MODEL_REPO_TYPE in {"model", "space"} else "model",
                )
            else:
                p = LABELS_FILE
            with open(p, "r", encoding="utf-8") as f:
                return [x.strip() for x in f if x.strip()]
    except Exception as e:
        print(f"[labels] ignorando erro: {e}")
    return []


def _init_once() -> Tuple[bool, str]:
    global _SERVING_FN, _LABELS, _LAST_INIT_ERROR
    if _SERVING_FN is not None:
        return True, "ok"
    try:
        pb_local = _download_model()
        sm_dir = _prepare_saved_model_dir(pb_local)
        model = tf.saved_model.load(sm_dir)
        # assinatura padrão esperada pelo Dancer Flow/Vertex TF Serving
        serving = model.signatures.get("serving_default")
        if serving is None:
            raise RuntimeError("SavedModel sem assinatura 'serving_default'.")
        _SERVING_FN = serving
        _LABELS = _maybe_labels()
        _LAST_INIT_ERROR = None
        return True, "ok"
    except Exception as e:
        _LAST_INIT_ERROR = f"{type(e).__name__}: {e}"
        return False, _LAST_INIT_ERROR


# =========================
# Pré/Pós-processamento
# =========================

def _preprocess_image_to_bytes(pil_img: Image.Image) -> bytes:
    img = pil_img.convert("RGB").resize((IMG_SIZE, IMG_SIZE))
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    return buf.getvalue()


def _postprocess(scores: np.ndarray, model_labels: List[str]) -> List[Dict[str, float]]:
    idxs = np.argsort(scores)[-TOPK:][::-1]
    out: List[Dict[str, float]] = []
    for i in idxs:
        label = model_labels[i] if i < len(model_labels) and model_labels[i] else (
            _LABELS[i] if i < len(_LABELS) else f"class_{i}"
        )
        out.append({"index": int(i), "label": label, "score": float(scores[i])})
    return out


# =========================
# Funções de UI
# =========================

def _signature_info() -> Dict[str, Dict[str, str]]:
    ok, err = _init_once()
    if not ok:
        return {"init_error": err}
    inputs = {k: str(v) for k, v in _SERVING_FN.structured_input_signature[1].items()}
    outputs = {k: str(v) for k, v in _SERVING_FN.structured_outputs.items()}
    return {"inputs": inputs, "outputs": outputs}


def _diagnostics() -> Dict[str, object]:
    ok, err = _init_once()
    return {
        "ok": ok,
        "error": err if not ok else None,
        "env": {
            "MODEL_URL": MODEL_URL or None,
            "MODEL_REPO": MODEL_REPO or None,
            "MODEL_REPO_TYPE": MODEL_REPO_TYPE,
            "MODEL_FILE": MODEL_FILE,
            "IMG_SIZE": IMG_SIZE,
            "TOPK": TOPK,
        },
    }


def infer(image: Image.Image):
    if image is None:
        raise ValueError("Envie uma imagem.")
    ok, err = _init_once()
    if not ok:
        raise RuntimeError(f"Modelo não inicializado: {err}")

    image_bytes = _preprocess_image_to_bytes(image)

    # Assinatura típica de TF Serving com bytes:
    #   inputs: image_bytes: tf.string, key: tf.string
    #   outputs: scores: tf.float32 [1, N], labels: tf.string [1, N]
    result = _SERVING_FN(
        image_bytes=tf.convert_to_tensor([image_bytes]),
        key=tf.convert_to_tensor(["0"]),
    )

    # Converte tensores nomeados
    scores = result.get("scores")
    labels = result.get("labels")
    if scores is None:
        raise KeyError("Saída 'scores' não encontrada na assinatura do modelo.")
    np_scores = scores.numpy()[0]

    model_labels: List[str] = []
    if labels is not None:
        model_labels = [x.decode("utf-8") for x in labels.numpy()[0]]

    return _postprocess(np_scores, model_labels)


# =========================
# Gradio UI
# =========================

demo = gr.Blocks(title="RaspagemTF - SavedModel (.pb)")
with demo:
    gr.Markdown("## RaspagemTF — Inferência (SavedModel .pb)")
    with gr.Row():
        img = gr.Image(type="pil", label="Imagem")
        res = gr.JSON(label="Top-K")
    btn = gr.Button("Rodar inferência")
    btn.click(fn=infer, inputs=img, outputs=res)

    with gr.Accordion("Diagnóstico", open=False):
        d_btn = gr.Button("Rodar diagnóstico")
        d_out = gr.JSON()
        d_btn.click(fn=_diagnostics, inputs=None, outputs=d_out)

    with gr.Accordion("Assinaturas do modelo", open=False):
        s_btn = gr.Button("Mostrar assinatura")
        s_out = gr.JSON()
        s_btn.click(fn=_signature_info, inputs=None, outputs=s_out)

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