import requests
from bs4 import BeautifulSoup

import pandas as pd

import gradio as gr
import json

from urllib import robotparser
from urllib.parse import urljoin

from collections import defaultdict

def extract_data(url, selected_elements):
    # Vérification robots.txt
    try:
        base_url = f"https://{url.split('//')[1].split('/')[0]}/robots.txt"
        rp = robotparser.RobotFileParser()
        rp.set_url(base_url)
        rp.read()
        if not rp.can_fetch("*", url):
            return json.dumps({"error": "Scraping interdit par robots.txt."}, indent=2, ensure_ascii=False)
    except:
        pass

    # Headers et requête
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/91.0"}
    try:
        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        return f"Erreur : {e.status_code if hasattr(e, 'status_code') else str(e)}"

    # Parsing HTML
    soup = BeautifulSoup(response.text, "html.parser")
    elements = [e.lower().replace(" ", "_") for e in selected_elements]  # Normalisation
    data = {}

    # Extraction des éléments
    try:
        if "titles" in elements:
            titles = {}
            for level in range(1, 7):
                tag = f"h{level}"
                found = [t.text.strip() for t in soup.find_all(tag)]
                if found:
                    titles[tag] = found
            data["titles"] = titles or {"h1": ["Aucun titre"]}

        if "hierarchical_titles" in elements:
            headings = soup.find_all(["h1", "h2", "h3", "h4", "h5", "h6"])
            hierarchy = defaultdict(list)
            current_level = {1: None, 2: None, 3: None, 4: None, 5: None, 6: None}

            for h in headings:
                hierarchy[h.name].append(h.text.strip())

            data["hierarchical_titles"] = dict(hierarchy) or {"h1": ["Aucun titre"]}
        if "paragraphs" in elements:
            data["paragraphs"] = [p.text.strip() for p in soup.find_all("p")] or ["Aucun paragraphe"]
        if "links" in elements:
            data["links"] = [urljoin(url, a["href"]) for a in soup.find_all("a", href=True)] or ["Aucun lien"]
        if "lists" in elements:
            lists = soup.find_all(["ul", "ol"])
            data["lists"] = [[li.text.strip() for li in lst.find_all("li")] for lst in lists] or ["Aucune liste"]

        # Extraction des tableaux
        """if "tables" in elements:
            tables = soup.find_all("table")
            table_data = []
            for table in tables:
                rows = []
                for row in table.find_all("tr"):
                    cols = [col.text.strip() for col in row.find_all(["td", "th"])]
                    rows.append(cols)
                table_data.append(rows)
            data["tables"] = table_data or ["Aucun tableau trouvé"]

        return json.dumps(data or "Aucun élément sélectionné.", indent=2, ensure_ascii=False)
    except Exception as e:
        return f"Erreur d'extraction : {str(e)}"""

        if "tables" in elements:
            try:
                tables = pd.read_html(url)  # Utilisation de pandas pour lire toutes les tables de l'URL
                # Conversion des DataFrames en JSON pour Gradio
                table_json = [table.to_json(orient="split") for table in tables]
                data["tables"] = table_json if table_json else ["Aucun tableau trouvé"]
            except Exception as e:
                data["tables"] = f"Erreur lors de l'extraction des tableaux : {str(e)}"

        return json.dumps(data or "Aucun élément sélectionné.", indent=2, ensure_ascii=False)
    except Exception as e:
        return f"Erreur d'extraction : {str(e)}"

# Interface Gradio avec boutons à sélection multiple
interface = gr.Interface(
    fn=extract_data,
    inputs=[
        gr.Textbox(label="URL", placeholder="https://example.com"),
        gr.CheckboxGroup(
            choices=["Titles", "Hierarchical Titles", "Paragraphs", "Links", "Lists", "Tables"],
            label="Éléments à extraire",
            value=["Titles"]
        )
    ],
    outputs=gr.JSON(label="Résultat"),  # Remplace Textbox par JSON
    title="Extracteur Web",
    description="Cochez les éléments à extraire d'une page web (respecte robots.txt)."
)

interface.launch()