from PIL import Image
import pytesseract
import spaces
import gradio as gr
import fitz  # PyMuPDF
import ocrmypdf
import tempfile
import os

def extract_text_markdown(doc):
    markdown_output = ""
    image_counter = 1

    for page in doc:
        blocks = page.get_text("dict")["blocks"]
        elements = []

        for b in blocks:
            y = b["bbox"][1]
            if b["type"] == 0:  # Texto
                for line in b["lines"]:
                    line_y = line["bbox"][1]
                    line_text = " ".join([span["text"] for span in line["spans"]]).strip()
                    if line_text:
                        elements.append((line_y, line_text))
            elif b["type"] == 1:  # Imagen
                elements.append((y, f"[imagen_{image_counter}]()"))
                image_counter += 1

        elements.sort(key=lambda x: x[0])

        previous_y = None
        for y, content in elements:
            if previous_y is not None and abs(y - previous_y) > 10:
                markdown_output += "\n"
            markdown_output += content + "\n"
            previous_y = y

        markdown_output += "\n---\n\n"

    return markdown_output.strip()

def needs_ocr(doc):
    text_length = sum(len(page.get_text().strip()) for page in doc)
    image_count = sum(len(page.get_images(full=True)) for page in doc)
    return text_length < 500 or image_count > 0

@spaces.GPU
def convert(pdf_file):
    doc = fitz.open(pdf_file)
    markdown_output = ""
    image_counter = 1

    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text("text").strip()

        if len(text) > 30:
            # Página con texto normal
            markdown_output += extract_text_markdown([page]) + "\n"
        else:
            # Página sin texto: usar OCR por imagen
            pix = page.get_pixmap(dpi=300)
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            ocr_text = pytesseract.image_to_string(img, lang="spa")
            markdown_output += ocr_text.strip() + "\n"

        markdown_output += "\n---\n\n"

    return markdown_output.strip(), {}

gr.Interface(
    fn=convert,
    inputs=[gr.File(label="Sube tu PDF", type="filepath")],
    outputs=[gr.Text(label="Markdown estructurado"), gr.JSON(label="Metadata")],
).launch()