import gradio as gr
from groq import Groq
import os
import base64
from io import BytesIO
import pandas as pd
import re
import tempfile
import unicodedata
from openpyxl import Workbook
from openpyxl.drawing.image import Image as OpenpyxlImage
from openpyxl.styles import Font, Alignment
from PIL import Image
import fitz  # PyMuPDF pour gérer les PDF

# Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

def image_to_base64(image):
    """Convert PIL image to base64 string for Groq API."""
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def pdf_to_images(pdf_path):
    """Convert PDF pages to PIL images."""
    images = []
    pdf_document = fitz.open(pdf_path)
    for page_num in range(len(pdf_document)):
        page = pdf_document[page_num]
        # Augmenter la résolution pour une meilleure qualité
        mat = fitz.Matrix(2.0, 2.0)  # zoom x2
        pix = page.get_pixmap(matrix=mat)
        img_data = pix.tobytes("jpeg")
        img = Image.open(BytesIO(img_data))
        images.append(img)
    pdf_document.close()
    return images

def parse_markdown_table_to_df(table_text):
    """Parse un tableau Markdown en Pandas DataFrame de manière robuste."""
    # Nettoyer les <br> en \n pour les sauts de ligne dans les cellules
    table_text = re.sub(r'<br>', '\n', table_text)

    lines = table_text.split('\n')

    # Ignorer les lignes vides, mais garder toutes les lignes non-séparateurs
    data_lines = []
    separator_pattern = r'\|[-| :]+\|'
    for line in lines:
        stripped = line.strip()
        if not stripped:
            continue
        if re.match(separator_pattern, stripped):
            continue
        data_lines.append(stripped)

    if len(data_lines) < 1:
        return pd.DataFrame({"Erreur": ["Aucun tableau Markdown trouvé dans la réponse"]})

    # Extraire les en-têtes (première ligne non vide)
    header_line = data_lines[0]
    headers = [h.strip() for h in header_line.split('|')[1:-1]]
    num_columns = len(headers)

    if num_columns == 0:
        return pd.DataFrame({"Erreur": ["Aucun en-tête valide trouvé"]})

    # Extraire les lignes de données (lignes suivantes)
    rows = []
    for line in data_lines[1:]:
        cells = line.split('|')[1:-1]
        cleaned_cells = [cell.strip() for cell in cells]
        # Gérer le mismatch de colonnes
        if len(cleaned_cells) < num_columns:
            cleaned_cells.extend([''] * (num_columns - len(cleaned_cells)))
        elif len(cleaned_cells) > num_columns:
            cleaned_cells = cleaned_cells[:num_columns]
        rows.append(cleaned_cells)

    # Créer le DataFrame
    df = pd.DataFrame(rows, columns=headers)
    # Filtrer les lignes entièrement vides
    df = df.loc[df.apply(lambda row: any(cell.strip() != '' for cell in row), axis=1)]
    return df if not df.empty else pd.DataFrame({"Erreur": ["Aucune donnée valide extraite"]})

def extract_filename_additional_and_table(response):
    """Extraire le nom de fichier, le texte additionnel et le tableau Markdown de la réponse structurée."""
    filename = "tableau_extrait"
    additional_text = ""
    table_text = ""

    if 'Filename:' in response:
        parts = response.split('Filename:', 1)[1].split('\n', 1)
        filename = parts[0].strip().replace('.xlsx', '')
        remaining = parts[1] if len(parts) > 1 else ""
    else:
        remaining = response

    if 'Additional text:' in remaining:
        parts = remaining.split('Additional text:', 1)[1].split('Table:', 1)
        additional_text = parts[0].strip()
        table_text = parts[1].strip() if len(parts) > 1 else ""
    else:
        table_text = remaining.strip()

    # Sanitizer le nom de fichier
    filename = ''.join(c for c in unicodedata.normalize('NFD', filename) if unicodedata.category(c) != 'Mn')
    filename = re.sub(r'[^a-zA-Z0-9_-]', '_', filename)
    filename = filename[:50]

    return filename, additional_text, table_text

def process_single_image(image, prompt):
    """Process a single image with the Groq API."""
    base64_image = image_to_base64(image)
    
    completion = client.chat.completions.create(
        model="meta-llama/llama-4-scout-17b-16e-instruct",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
                    }
                ]
            }
        ],
        temperature=0.2,
        max_completion_tokens=4096,
        top_p=1,
        stream=False,
        stop=None
    )
    
    return completion.choices[0].message.content.strip()

def create_excel_file(filename, df, additional_text, images):
    """Create Excel file with table, additional info, and original images."""
    excel_path = f"/tmp/{filename if filename else 'tableau_extrait'}.xlsx"
    wb = Workbook()

    # Feuille pour le tableau
    ws_table = wb.active
    ws_table.title = 'Tableau_Extrait'
    if not df.empty and "Erreur" not in df.columns:
        # Écrire les en-têtes en gras
        for col_num, value in enumerate(df.columns.values, start=1):
            cell = ws_table.cell(row=1, column=col_num, value=value)
            cell.font = Font(bold=True)
            cell.alignment = Alignment(wrap_text=True, vertical='top')
            ws_table.column_dimensions[chr(64 + col_num)].width = max(len(str(value)) * 1.2, 15)

        # Écrire les données avec wrap text
        for row_num, row in enumerate(df.values, start=2):
            for col_num, value in enumerate(row, start=1):
                cell = ws_table.cell(row=row_num, column=col_num, value=value)
                cell.alignment = Alignment(wrap_text=True, vertical='top')
                ws_table.row_dimensions[row_num].height = max(len(str(value).split('\n')) * 15, 20)
    else:
        ws_table.cell(row=1, column=1, value="Erreur lors du parsing du tableau")
        ws_table.column_dimensions['A'].width = 50
        ws_table.row_dimensions[1].height = 20

    # Feuille pour le texte additionnel
    if additional_text:
        ws_additional = wb.create_sheet(title='Infos_Supplementaires')
        lines = additional_text.split('\n')
        for row_num, line in enumerate(lines, start=1):
            cell = ws_additional.cell(row=row_num, column=1, value=line)
            cell.alignment = Alignment(wrap_text=True, vertical='top')
            ws_additional.column_dimensions['A'].width = max(len(line) * 1.2, 50)
            ws_additional.row_dimensions[row_num].height = max(len(line.split('\n')) * 15, 20)

    # Feuilles pour les images originales
    for idx, img in enumerate(images, start=1):
        img_path = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg").name
        img.save(img_path, format="JPEG")
        
        sheet_name = f'Image_Page_{idx}' if len(images) > 1 else 'Image_Originale'
        ws_image = wb.create_sheet(title=sheet_name)
        excel_img = OpenpyxlImage(img_path)
        ws_image.add_image(excel_img, 'A1')
        
        os.unlink(img_path)

    wb.save(excel_path)
    return excel_path

def process_file(file):
    """Process uploaded file (image or PDF) and generate Excel."""
    if file is None:
        return "Veuillez uploader un fichier (image ou PDF).", None

    try:
        file_path = file.name if hasattr(file, 'name') else file
        file_extension = os.path.splitext(file_path)[1].lower()
        
        # Déterminer si c'est un PDF ou une image
        images = []
        if file_extension == '.pdf':
            images = pdf_to_images(file_path)
            if not images:
                return "Erreur : Le PDF ne contient aucune page.", None
        else:
            # C'est une image
            if isinstance(file, str):
                images = [Image.open(file)]
            else:
                images = [file]

        # Prompt optimisé
        prompt = (
            "Analyse l'image et extrait tout le contenu avec précision. "
            "D'abord, suggère un nom de fichier descriptif et précis pour l'Excel basé sur le contenu principal de l'image "
            "(inclue des éléments clés comme le mois, l'année, ou des identifiants uniques, ex: 'Registre_Heures_Juin_2016'). "
            "Ensuite, extrait TOUT texte additionnel autour, au-dessus, en-dessous ou à côté du tableau "
            "(titres, en-têtes de page, notes, pieds de page, logos, etc.), en le recopiant mot pour mot, même si c'est dispersé. "
            "Si aucun texte additionnel, laisse vide. "
            "Enfin, extrait le tableau COMPLET en entier, en recopiant TOUTES les lignes et colonnes à l'identique, "
            "y compris les lignes vides ou partielles si elles existent dans l'image. "
            "Les en-têtes du tableau (première ligne avec les titres des colonnes) doivent TOUJOURS être inclus dans la section Table, "
            "JAMAIS dans Additional text. "
            "Utilise un format Markdown pour le tableau avec des | pour les colonnes et une ligne |---|---| pour les séparateurs. "
            "Assure-toi que CHAQUE ligne (en-têtes, séparateurs, données) a EXACTEMENT le même nombre de colonnes. "
            "Pour les sauts de ligne dans une cellule, utilise \\n au lieu de <br>. "
            "Remplis les cellules vides avec '' si nécessaire pour maintenir l'alignement. "
            "N'ajoute aucun texte explicatif dans le tableau. "
            "Structure ta réponse exactement comme suit :\n"
            "Filename: [nom_suggéré_précis]\n"
            "Additional text: [tout le texte additionnel, séparé par \\n si plusieurs lignes ; sinon vide]\n"
            "Table:\n"
            "[le tableau Markdown ici]"
        )

        # Traiter chaque page/image
        all_responses = []
        all_tables = []
        all_additional_texts = []
        filename = "tableau_extrait"
        
        for idx, img in enumerate(images):
            response = process_single_image(img, prompt)
            all_responses.append(f"--- Page {idx + 1} ---\n{response}")
            
            # Extraire les informations
            current_filename, additional_text, table_text = extract_filename_additional_and_table(response)
            
            # Utiliser le nom de fichier de la première page
            if idx == 0 and current_filename:
                filename = current_filename
            
            if additional_text:
                all_additional_texts.append(f"Page {idx + 1}:\n{additional_text}")
            
            # Parser le tableau
            df = parse_markdown_table_to_df(table_text)
            all_tables.append((idx + 1, df))

        # Combiner tous les tableaux en un seul DataFrame
        combined_df = pd.DataFrame()
        for page_num, df in all_tables:
            if not df.empty and "Erreur" not in df.columns:
                # Ajouter une colonne "Page" si plusieurs pages
                if len(all_tables) > 1:
                    df.insert(0, 'Page', page_num)
                combined_df = pd.concat([combined_df, df], ignore_index=True)

        # Si aucun tableau valide n'a été trouvé, utiliser le premier DataFrame d'erreur
        if combined_df.empty:
            combined_df = all_tables[0][1] if all_tables else pd.DataFrame({"Erreur": ["Aucune donnée extraite"]})

        # Combiner tous les textes additionnels
        combined_additional_text = "\n\n".join(all_additional_texts) if all_additional_texts else ""

        # Créer le fichier Excel
        excel_path = create_excel_file(filename, combined_df, combined_additional_text, images)

        # Créer la réponse complète
        full_response = "\n\n".join(all_responses) if len(images) > 1 else all_responses[0]

        return full_response, excel_path

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

# Define Gradio interface
iface = gr.Interface(
    fn=process_file,
    inputs=gr.File(
        label="Uploader une image ou un PDF contenant un tableau",
        file_types=["image", ".pdf"]
    ),
    outputs=[
        gr.Textbox(
            label="Réponse de l'IA (texte additionnel + tableau Markdown)",
            lines=15
        ),
        gr.File(label="Télécharger le fichier Excel (avec tableau, infos supp. et images originales)")
    ],
    title="Extraction de Tableau depuis Image/PDF avec Groq et Export Excel",
    description=(
        "Uploader une image ou un PDF contenant un tableau. "
        "L'IA extrait le texte additionnel et le tableau, puis génère un Excel avec des feuilles séparées, "
        "y compris les images originales. Pour les PDF multi-pages, tous les tableaux sont combinés. "
        "Le résultat n'est pas parfait, veuillez vous relire pour vérifier l'exactitude des réponses. "
        "Les données sont privées et ne sont pas sauvegardées."
    ),
    examples=[
        ["exemple_tableau.jpg"],
        ["exemple_document.pdf"]
    ]
)

# Launch the interface
if __name__ == "__main__":
    iface.launch()