main.py

import os
import time
import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.action_chains import ActionChains
from webdriver_manager.chrome import ChromeDriverManager

def wait_for_download_to_complete(download_directory, timeout=15):
    """Vérifie que le téléchargement est bien terminé."""
    end_time = time.time() + timeout
    while time.time() < end_time:
        files = os.listdir(download_directory)
        if any(f.endswith('.crdownload') or f.endswith('.tmp') for f in files):
            time.sleep(0.2)
        else:
            return True
    return False

def lancer_telechargement(nom_produit, mots_cles, nomenclature_str, progress=gr.Progress()):
    """Fonction principale appelée par le bouton Gradio."""
    download_dir = os.path.join(os.getcwd(), "FDES_Telecharges")
    if not os.path.exists(download_dir):
        os.makedirs(download_dir)

    chrome_options = Options()
    chrome_options.page_load_strategy = 'eager'
    chrome_options.add_argument("--headless=new")  # Mode invisible
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--window-size=1920,1080")

    chrome_options.add_experimental_option(
        'prefs', {
            "download.default_directory": download_dir,
            "download.prompt_for_download": False,
            "download.directory_upgrade": True,
            "plugins.always_open_pdf_externally": True,
            "profile.default_content_settings.popups": 0,
            "profile.managed_default_content_settings.images": 2
        }
    )

    progress(0, desc="Démarrage du navigateur...")
    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)

    # Autoriser les téléchargements en mode Headless
    driver.command_executor._commands["send_command"] = ("POST", '/session/$sessionId/chromium/send_command')
    params = {'cmd': 'Page.setDownloadBehavior', 'params': {'behavior': 'allow', 'downloadPath': download_dir}}
    driver.execute("send_command", params)

    wait = WebDriverWait(driver, 10)
    actions = ActionChains(driver)

    def navigate_angular_node(keyword, expand=True):
        progress(0, desc=f"Recherche de '{keyword}'...")
        target_node = None
        # Astuce : Gérer les exclusions via une syntaxe spéciale (ex: "Chauffage |NOT| Ventilation")
        exclude_keyword = None
        if "|NOT|" in keyword:
            keyword, exclude_keyword = keyword.split("|NOT|")
            keyword = keyword.strip()
            exclude_keyword = exclude_keyword.strip()

        for _ in range(20):
            wrappers = driver.find_elements(By.CSS_SELECTOR, ".node-wrapper")
            for w in wrappers:
                text_lower = w.text.lower()
                # Si le mot-clé est trouvé
                if keyword.lower() in text_lower:
                    # On vérifie qu'il ne contient PAS le mot exclu (s'il y en a un)
                    if exclude_keyword and exclude_keyword.lower() in text_lower:
                        continue
                    # On vérifie que le noeud est bien affiché à l'écran
                    if w.is_displayed():
                        target_node = w
                        break
            if target_node:
                break
            time.sleep(0.2)

        if not target_node:
            raise Exception(f"❌ Dossier contenant '{keyword}' introuvable à l'écran.")

        driver.execute_script("arguments[0].scrollIntoView({behavior: 'auto', block: 'center'});", target_node)

        if expand:
            try:
                expander = target_node.find_element(By.CSS_SELECTOR, ".toggle-children-wrapper-collapsed")
                driver.execute_script("arguments[0].click();", expander)
                time.sleep(0.5)  # Laisser le temps au dossier de s'ouvrir
            except:
                pass  # Déjà ouvert ou pas de flèche
        else:
            content = target_node.find_element(By.CSS_SELECTOR, ".node-content-wrapper")
            driver.execute_script("arguments[0].click();", content)

    try:
        progress(0.1, desc="Ouverture du site INIES...")
        driver.get("https://base-inies.fr/consultation/recherche")

        wait.until(EC.element_to_be_clickable((By.CSS_SELECTOR, "button.nomenclature-select-btn")))

        # Accepter les cookies
        try:
            cookie_btn = driver.find_element(By.XPATH, "//*[contains(text(), 'compris') or contains(text(), 'Accepter')]")
            driver.execute_script("arguments[0].click();", cookie_btn)
        except:
            pass

        # 1. Remplir le champ "Nom de produit"
        if nom_produit:
            input_nom = driver.find_element(By.CSS_SELECTOR, 'input[formcontrolname="nomProduit"]')
            input_nom.send_keys(nom_produit)

        # 2. Remplir le champ "Mots clés"
        if mots_cles:
            input_mots = driver.find_element(By.CSS_SELECTOR, 'input[formcontrolname="motsCles"]')
            input_mots.send_keys(mots_cles)

        # 3. Gérer la nomenclature (si renseignée)
        if nomenclature_str:
            btn_nomenclature = driver.find_element(By.CSS_SELECTOR, "button.nomenclature-select-btn")
            driver.execute_script("arguments[0].click();", btn_nomenclature)
            wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, ".node-wrapper")))

            # Découper le chemin fourni par l'utilisateur (séparé par des ;)
            etapes = [e.strip() for e in nomenclature_str.split(';') if e.strip()]
            for i, etape in enumerate(etapes):
                est_dernier = (i == len(etapes) - 1)
                # Déplier tous les dossiers sauf le dernier qu'on doit sélectionner
                navigate_angular_node(etape, expand=not est_dernier)

            btn_ok = wait.until(EC.element_to_be_clickable((By.XPATH, "//div[contains(@class, 'app-modal-footer')]//button[contains(text(), 'Ok')]")))
            driver.execute_script("arguments[0].click();", btn_ok)
            time.sleep(0.5)

        # 4. Lancer la recherche globale
        progress(0.2, desc="Lancement de la recherche...")
        btn_rechercher = wait.until(EC.element_to_be_clickable((By.XPATH, "//button[contains(., 'Rechercher')]")))
        driver.execute_script("arguments[0].click();", btn_rechercher)

        # 5. Attente des résultats
        try:
            wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, ".result-panel .tree-node-leaf")))
            time.sleep(1)
        except:
            return "Aucun produit trouvé avec ces critères."

        leaf_nodes = driver.find_elements(By.CSS_SELECTOR, ".result-panel .tree-node-leaf .node-content-wrapper")
        num_products = len(leaf_nodes)

        if num_products == 0:
            return "Aucun produit trouvé avec ces critères."

        # 6. Boucle de téléchargement
        telecharges = 0
        for i in range(num_products):
            progress(0.2 + (0.8 * (i / num_products)), desc=f"Téléchargement: {i + 1}/{num_products} produits...")

            wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, ".result-panel .tree-node-leaf")))
            leaf_nodes = driver.find_elements(By.CSS_SELECTOR, ".result-panel .tree-node-leaf .node-content-wrapper")

            if i >= len(leaf_nodes):
                break

            product_to_click = leaf_nodes[i]

            try:
                driver.execute_script("arguments[0].scrollIntoView({behavior: 'auto', block: 'center'});", product_to_click)
                actions.move_to_element(product_to_click).click().perform()
            except:
                continue

            # Onglet Documents et téléchargement
            try:
                wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, ".produit-details")))

                tab_docs = wait.until(EC.element_to_be_clickable((By.XPATH, "//button[contains(@class, 'tablinks') and contains(., 'Documents')]")))
                driver.execute_script("arguments[0].click();", tab_docs)

                xpath_pdf = "//td[contains(., 'FDES/PEP')]/following-sibling::td/a"
                pdf_link = wait.until(EC.presence_of_element_located((By.XPATH, xpath_pdf)))

                driver.execute_script("arguments[0].scrollIntoView({behavior: 'auto', block: 'center'});", pdf_link)
                driver.execute_script("arguments[0].click();", pdf_link)

                # Attendre la fin du fichier .crdownload
                wait_for_download_to_complete(download_dir)
                telecharges += 1
            except:
                pass  # Pas de FDES pour ce produit

            driver.back()
            wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, ".result-panel .tree-node-leaf")))

        return f"✅ Opération terminée avec succès !\n\n{telecharges} fichier(s) FDES/PEP téléchargé(s) sur un total de {num_products} produit(s) trouvé(s).\nDossier : {download_dir}"

    except Exception as e:
        return f"❌ Une erreur est survenue : {str(e)}"

    finally:
        driver.quit()

# ---------------------------------------------------------
# INTERFACE GRAPHIQUE AVEC GRADIO
# ---------------------------------------------------------
with gr.Blocks(title="Extracteur INIES", theme=gr.themes.Soft()) as interface:
    gr.Markdown("# 🏢 Extracteur Automatique - Base INIES")
    gr.Markdown("Renseignez un ou plusieurs champs. Le robot s'occupera de piloter le site en arrière-plan pour récupérer les PDFs (FDES/PEP) correspondants.")

    with gr.Row():
        with gr.Column(scale=2):
            nom_produit = gr.Textbox(label="Nom de produit", placeholder="Ex: AQUACIAT POWER")
            mots_cles = gr.Textbox(label="Mot(s) clé(s)", placeholder="Ex: Pompe à chaleur")

            gr.Markdown("**Famille de produit** (Nomenclature)\nSéparez chaque étape par un point-virgule (`;`).\n*Astuce : Utilisez `|NOT|` pour exclure un terme si plusieurs dossiers se ressemblent.*")
            nomenclature = gr.Textbox(
                label="",
                placeholder="Ex: électriques; génie climatique |NOT| électriques; Chauffage |NOT| Ventilation; Pompes à chaleur collectives; Air/eau",
                value="électriques; génie climatique |NOT| électriques; Chauffage |NOT| Ventilation; Pompes à chaleur collectives; Air/eau",
                lines=2
            )

            btn_lancer = gr.Button("🚀 Lancer la recherche et le téléchargement", variant="primary")

        with gr.Column(scale=1):
            resultat = gr.Textbox(label="Statut de l'opération", lines=8, interactive=False)

    # Lier le bouton à la fonction Selenium
    btn_lancer.click(
        fn=lancer_telechargement,
        inputs=[nom_produit, mots_cles, nomenclature],
        outputs=resultat
    )

if __name__ == "__main__":
    # Lancement de l'interface
    print("Démarrage de l'interface Gradio...")
    interface.launch()