import os
import gradio as gr
import base64
import io
import re
from PIL import Image
import fitz  # PyMuPDF

# --- CONFIGURATION ---
# Récupère la clé HF depuis les secrets du Space (sécurisé)
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError(
        "❌ Token HF manquant. Ajoutez HF_TOKEN dans les 'Repository Secrets' de votre Space."
    )

MODEL_ID  = "google/gemma-4-E4B-it"
API_URL   = f"https://api-inference.huggingface.co/models/{MODEL_ID}/v1/chat/completions"

# --- FONCTIONS AUXILIAIRES ---

def encode_image(image: Image.Image) -> str:
    """Convertit une image PIL en base64 JPEG."""
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG", quality=85)
    return base64.b64encode(buffer.getvalue()).decode()


def pdf_to_image(pdf_path: str) -> Image.Image:
    """Convertit la première page d'un PDF en image PIL."""
    doc = fitz.open(pdf_path)
    page = doc.load_page(0)
    mat  = fitz.Matrix(150 / 72, 150 / 72)   # 150 DPI — lisible et léger
    pix  = page.get_pixmap(matrix=mat)
    img  = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    doc.close()
    return img


def sanitize_filename(name: str) -> str:
    """Nettoie le nom retourné par le modèle pour en faire un nom de fichier valide."""
    name = name.strip().strip('"').strip("'")
    name = re.sub(r"[^\w\s\-]", "", name)
    name = re.sub(r"\s+", "_", name)
    return name[:80] if name else "document_sans_nom"


# --- FONCTION PRINCIPALE ---

def generate_filename(file, custom_prompt: str):
    """
    Génère un nom de fichier à partir d'une image ou d'un PDF (première page).
    Args:
        file     : Fichier uploadé (image ou PDF).
        custom_prompt : Prompt/format personnalisé pour le nom.
    Returns:
        tuple: (aperçu PIL Image, nom généré str)
    """
    # 1. Vérification du fichier
    if file is None:
        return None, "❌ Veuillez uploader un fichier (image ou PDF)."

    try:
        # 2. Traitement selon le type de fichier
        if file.name.lower().endswith(".pdf"):
            image = pdf_to_image(file.name)
        else:
            image = Image.open(file.name).convert("RGB")

        # 3. Encodage base64
        image_base64 = encode_image(image)

        # 4. Prompt
        if custom_prompt and custom_prompt.strip():
            prompt = custom_prompt.strip()
        else:
            prompt = (
                "Analyse ce document. Génère un nom de fichier court et descriptif "
                "en snake_case (3 à 6 mots max).\n"
                "Réponds UNIQUEMENT par le nom de fichier, sans extension, sans explication."
            )

        # 5. Appel à l'Inference API HF (même structure que ton ancien script Mistral)
        import requests
        response = requests.post(
            API_URL,
            headers={
                "Authorization": f"Bearer {HF_TOKEN}",
                "Content-Type": "application/json",
            },
            json={
                "model": MODEL_ID,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                },
                            },
                        ],
                    }
                ],
                "max_tokens": 60,
                "temperature": 0.1,
            },
            timeout=60,
        )

        # 6. Traitement de la réponse
        if response.status_code == 200:
            raw   = response.json()["choices"][0]["message"]["content"].strip()
            clean = sanitize_filename(raw)
            return image, f"{clean}.pdf"
        else:
            return image, f"❌ Erreur API HF : {response.status_code} — {response.text}"

    except Exception as e:
        return None, f"❌ Erreur : {str(e)}"


# --- INTERFACE GRADIO ---

with gr.Blocks(title="PDF Auto-Namer") as demo:
    gr.Markdown("""
    # 📄 PDF Auto-Namer
    **Basé sur Gemma 4 E4B (multimodal)** via Hugging Face Inference API
    Upload un **fichier image ou PDF** → la première page est analysée → un nom est généré.
    """)

    with gr.Row():
        with gr.Column():
            file_input = gr.File(
                label="Document (Image ou PDF)",
                file_types=[".jpg", ".jpeg", ".png", ".pdf"],
                height=200,
            )
            prompt_input = gr.Textbox(
                label="Format personnalisé (optionnel)",
                placeholder='Ex: "Format: Contrat_[TYPE]_[CLIENT]_[JJMMAAAA]"',
                lines=3,
            )
            btn = gr.Button("🚀 Générer le nom", variant="primary")

        with gr.Column():
            preview = gr.Image(label="👁️ Aperçu — Page 1", type="pil")
            output  = gr.Textbox(label="📝 Nom de fichier généré", lines=2)

    btn.click(
        fn=generate_filename,
        inputs=[file_input, prompt_input],
        outputs=[preview, output],
    )

    gr.Markdown(
        "---\n"
        "ℹ️ *Modèle : `google/gemma-4-E4B-it` · "
        "La clé `HF_TOKEN` doit être définie dans les Secrets du Space.*"
    )

demo.launch(show_error=True)