import fitz
import pytesseract
from PIL import Image
from io import BytesIO
import shutil

# O path do tesseract será configurado dinamicamente dentro da função para evitar 
# problemas caso ele seja instalado enquanto o servidor já estiver rodando.

def make_pdf_searchable(pdf_bytes: bytes, lang: str = "por") -> bytes:
    """
    Recebe um PDF digitalizado, converte cada página em imagem, passa pelo Tesseract OCR
    para extrair o texto e gerar um novo PDF pesquisável (onde o texto fica invisível 
    sobreposto à imagem perfeitamente).
    """
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    new_doc = fitz.open()

    import os
    if shutil.which("tesseract") is None:
        if os.path.exists(r"C:\Program Files\Tesseract-OCR\tesseract.exe"):
            pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
        elif os.path.exists(r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe"):
            pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe"

    for page_num in range(len(doc)):
        page = doc[page_num]
        
        # Converte a página para imagem (resolução otimizada para evitar timeout em PDFs grandes)
        pix = page.get_pixmap(dpi=150, alpha=False)
        
        # Converte o Pixmap do PyMuPDF para PIL Image
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        
        # Envia a imagem para o Tesseract, pedindo a saída diretamente como PDF pesquisável
        try:
            import os, platform
            # Configura para procurar os idiomas no diretório atual do backend apenas no Windows (local)
            if platform.system() == "Windows":
                os.environ["TESSDATA_PREFIX"] = os.path.dirname(__file__)
            
            pdf_page_bytes = pytesseract.image_to_pdf_or_hocr(
                img, 
                extension='pdf', 
                lang=lang
            )
            
            # Abre o mini-PDF gerado pelo Tesseract
            page_doc = fitz.open(stream=pdf_page_bytes, filetype="pdf")
            
            # Adiciona a página ao nosso documento final
            new_doc.insert_pdf(page_doc)
            page_doc.close()
            
            # Libera memória no final do loop
            del pix, img, pdf_page_bytes
            
        except pytesseract.TesseractNotFoundError:
            raise Exception("Tesseract-OCR não encontrado. Por favor, instale o Tesseract no servidor.")
        except Exception as e:
            raise Exception(f"Erro no OCR da página {page_num + 1}: {str(e)}")

    doc.close()
    
    out = BytesIO()
    # Salva o novo PDF
    new_doc.save(out, garbage=4, deflate=True)
    new_doc.close()
    
    return out.getvalue()