diarios_oficiais_search_alems/diarios/models.py

import os
from urllib.parse import urlparse
import requests
from babel.dates import format_date
from django.core.files.base import ContentFile
from django.db import models
from django.core.exceptions import ValidationError
import PyPDF2
import pdfplumber
from asgiref.sync import async_to_sync
import fitz


class TipoDiarioOficial(models.Model):
    """Representa um tipo de Diário Oficial (e.g., Municipal, Estadual, Federal)."""

    nome = models.CharField(max_length=100, unique=True)

    def __str__(self):
        """Retorna o nome do tipo de Diário Oficial."""
        return self.nome

    class Meta:
        verbose_name_plural = "Tipos de Diários Oficiais"


class DiarioOficial(models.Model):
    """Modelo que representa um Diário Oficial, contendo data, arquivo PDF, tipo e link.

    Attributes:
        data (DateField): Data de publicação do Diário Oficial.
        arquivo (FileField): Arquivo PDF do Diário Oficial (opcional).
        tipo (ForeignKey): Tipo do Diário Oficial (Municipal, Estadual, etc.).
        numero (CharField): Número de identificação único do Diário.
        link (URLField): URL para o Diário Oficial (opcional).
    """

    data = models.DateField()
    arquivo = models.FileField(upload_to="diarios_oficiais/", blank=True, null=True)
    tipo = models.ForeignKey(
        TipoDiarioOficial,
        blank=True,
        null=True,
        on_delete=models.SET_NULL,
        related_name="diarios",
    )
    numero = models.CharField(max_length=20, unique=True)
    link = models.URLField(blank=True, null=True, unique=True)

    def save(self, *args, **kwargs):
        """Salva o Diário Oficial, baixa o PDF (se houver link) e extrai páginas."""
        super().save(*args, **kwargs)

        if self.link and not self.arquivo:
            self._download_pdf_from_link()

        if self.arquivo and not self.paginas:
            self._extract_pdf_pages()

        super().save(*args, **kwargs)

    def clean(self):
        """Valida o modelo antes de salvar (chamado automaticamente no admin/form)."""
        super().clean()
        if not self.arquivo and not self.link:
            raise ValidationError("Informe um arquivo ou um link para o Diário.")

    def _validar_link(self):
        """Verifica se o link é um PDF válido."""
        if not self.link.lower().endswith(".pdf"):
            raise ValidationError("O link deve apontar para um arquivo PDF.")

    def _download_pdf_from_link(self):
        """Faz download do PDF a partir do link e salva no campo `arquivo`.

        Raises:
            ValidationError: Se o download falhar.
        """
        try:
            response = requests.get(self.link)
            response.raise_for_status()

            parsed_url = urlparse(self.link)
            file_name = os.path.basename(parsed_url.path) or f"diario_{self.numero}.pdf"

            self.arquivo.save(file_name, ContentFile(response.content), save=True)
        except requests.RequestException as e:
            raise ValidationError(f"Não foi possível baixar o PDF: {e}")

    
    def _extract_pdf_pages(self):
        """Extrai o texto de cada página do PDF usando PyMuPDF."""
        try:
            with self.arquivo.open("rb") as pdf_file:
                pdf_document = fitz.open(stream=pdf_file.read(), filetype="pdf")
                self.paginas.all().delete()

                for i, page in enumerate(pdf_document):
                    try:
                        # Extração simples
                        page_text = page.get_text("text")
                        
                        # Você pode experimentar com outros métodos:
                        # page.get_text("blocks")
                        # page.get_text("words")
                        # para melhor tratamento de duas colunas

                        if page_text and page_text.strip():
                            PageDiarioOficial.objects.create(
                                diario=self,
                                numero=i + 1,
                                conteudo=page_text.strip(),
                            )
                    except Exception as page_error:
                        PageDiarioOficial.objects.create(
                            diario=self,
                            numero=i + 1,
                            conteudo=f"[Erro na extração do texto: {str(page_error)}]",
                        )

        except Exception as pdf_error:
            raise ValidationError(f"Não foi possível processar o PDF: {pdf_error}")

    def _process_pdf_pages(self, pdf):
        """Processa cada página do PDF e salva seu conteúdo.

        Args:
            pdf (PdfReader): Objeto PDF carregado.
        """
        self.paginas.all().delete()

        for i, pagina in enumerate(pdf.pages):
            try:
                page_text = pagina.extract_text()
                if page_text and page_text.strip():
                    PageDiarioOficial.objects.create(
                        diario=self,
                        numero=i + 1,
                        conteudo=page_text.strip(),
                    )
            except Exception as page_error:
                PageDiarioOficial.objects.create(
                    diario=self,
                    numero=i + 1,
                    conteudo=f"[Erro na extração do texto: {str(page_error)}]",
                )
                continue

    @property
    def data_formatada(self):
        """Retorna a data formatada em português (e.g., '1 de Janeiro de 2023')."""
        return format_date(self.data, format="long", locale="pt_BR")

    @property
    def is_online(self):
        """Verifica se o Diário possui um link (online)."""
        return bool(self.link)

    def __str__(self):
        """Representação em string do Diário Oficial."""
        tipo_nome = self.tipo.nome if self.tipo else "Sem Tipo"
        return f"Diário {tipo_nome} nº {self.numero}, {self.data_formatada}"

    class Meta:
        constraints = [models.UniqueConstraint(fields=["numero"], name="unique_numero")]
        verbose_name_plural = "Diários Oficiais"


class PageDiarioOficial(models.Model):
    """Representa uma página de um Diário Oficial com seu conteúdo textual.

    Attributes:
        diario (ForeignKey): Diário Oficial associado.
        layout_duas_colunas (BooleanField): Indica se a página tem duas colunas.
        numero (PositiveIntegerField): Número da página no Diário.
        conteudo (TextField): Texto extraído da página.
    """

    diario = models.ForeignKey(
        DiarioOficial, on_delete=models.CASCADE, related_name="paginas"
    )
    layout_duas_colunas = models.BooleanField(default=False)
    numero = models.PositiveIntegerField()
    conteudo = models.TextField()

    class Meta:
        unique_together = ("diario", "numero")
        verbose_name = "Página de Diário Oficial"
        verbose_name_plural = "Páginas de Diários Oficiais"

    def __str__(self):
        """Representação em string da página (e.g., 'Página 1 do Diário 123')."""
        return f"Página {self.numero} do Diário {self.diario.numero}"

    def save(self, *args, **kwargs):
        from diarios.documents import DiarioOficialDocument

        super().save(*args, **kwargs)
        try:
            print(f"Reindexando diario {self.diario}")
            DiarioOficialDocument().update(self.diario)
        except Exception as e:
            print(f"Erro ao reindexar DiarioOficial {self.diario.id} (pai da página {self.id}) no Elasticsearch: {e}")