import os
import shutil
import gradio as gr
from PIL import Image
import tempfile
from pathlib import Path

MODEL_DIR = os.path.join(os.path.dirname(__file__), "models", "decimer")
os.makedirs(MODEL_DIR, exist_ok=True)
os.environ["DECIMER_CACHE_DIR"] = MODEL_DIR

def ensure_model_in_models():
    """Vérifie si le modèle est bien dans /models/decimer, sinon copie depuis le cache utilisateur."""
    # Emplacement typique du cache DECIMER (à adapter si besoin)
    home = Path.home()
    possible_paths = [
        home / ".cache" / "decimer",
        home / ".data" / "DECIMER-V2"
    ]
    # Test si déjà dans /models/decimer
    if any(Path(MODEL_DIR).glob("*")):
        return  # déjà présent, rien à faire

    # Sinon, on force le téléchargement dans le cache via un appel bidon
    from DECIMER import predict_SMILES
    img_path = os.path.join(MODEL_DIR, "dummy.png")
    with open(img_path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\xdacd\xf8\x0f\x00\x01\x05\x01\x02\xa2Z\xb2\x00\x00\x00\x00IEND\xaeB`\x82")
    try:
        predict_SMILES(img_path)
    except Exception:
        pass

    # On cherche le modèle dans le cache
    for p in possible_paths:
        if p.exists():
            # Recopie tous les fichiers/dossiers du cache dans /models/decimer
            for src in p.iterdir():
                dest = Path(MODEL_DIR) / src.name
                if src.is_dir():
                    if dest.exists():
                        shutil.rmtree(dest)
                    shutil.copytree(src, dest)
                else:
                    shutil.copy2(src, dest)
            print(f"Modèles copiés de {p} vers {MODEL_DIR}")

ensure_model_in_models()

from DECIMER import predict_SMILES

def predict_fn(img):
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        img.save(tmp, format="PNG")
        path = tmp.name
    return predict_SMILES(path)

iface = gr.Interface(
    fn=predict_fn,
    inputs=gr.Image(type="pil"),
    outputs="text"
)

iface.launch(share=True)