import os
import time
from django.core.files import File
from django.core.management.base import BaseCommand
from diarios.models import DiarioOficial
from diarios.signals import update_document, delete_document
from django.db.models.signals import post_save, post_delete


class Command(BaseCommand):
    help = "Importa arquivos PDF de diários oficiais e associa aos objetos existentes no banco."

    def add_arguments(self, parser):
        parser.add_argument(
            "pasta",
            type=str,
            help="Caminho para a pasta contendo os arquivos PDF (nomes devem conter o número do diário)",
        )

    def handle(self, *args, **options):
        post_save.disconnect(update_document, sender=DiarioOficial)
        post_delete.disconnect(delete_document, sender=DiarioOficial)

        pasta = options["pasta"]
        if not os.path.isdir(pasta):
            self.stderr.write(
                self.style.ERROR(f"A pasta fornecida não existe: {pasta}")
            )
            return

        arquivos_pdf = [
            f
            for f in os.listdir(pasta)
            if f.lower().endswith(".pdf") and "Diário_Oficial_Eletrônico_nº_" in f
        ]

        if not arquivos_pdf:
            self.stdout.write(
                self.style.WARNING("Nenhum arquivo PDF válido encontrado na pasta.")
            )
            return

        total = len(arquivos_pdf)
        erros = []
        atualizados = []
        start_time = time.time()

        # Mapeia os números aos nomes de arquivos
        numero_para_arquivo = {
            f.replace("Diário_Oficial_Eletrônico_nº_", "")
            .replace(".pdf", "")
            .strip("_-"): f
            for f in arquivos_pdf
        }

        # Busca todos os objetos de uma vez
        diarios_existentes = DiarioOficial.objects.in_bulk(
            numero_para_arquivo.keys(), field_name="numero"
        )

        for idx, (numero, nome_arquivo) in enumerate(
            numero_para_arquivo.items(), start=1
        ):
            try:
                diario = diarios_existentes.get(numero)
                if not diario:
                    raise DiarioOficial.DoesNotExist()

                caminho_pdf = os.path.join(pasta, nome_arquivo)
                with open(caminho_pdf, "rb") as f:
                    diario.arquivo.save(nome_arquivo, File(f), save=True)

                atualizados.append(diario.pk)
                elapsed = time.time() - start_time
                remaining = (elapsed / idx) * (total - idx)
                self.stdout.write(
                    f"[{idx}/{total}] Atualizado: {nome_arquivo} | Estimativa restante: {remaining:.1f}s"
                )

            except DiarioOficial.DoesNotExist:
                msg = f"Não encontrado no banco: {nome_arquivo}"
                erros.append(msg)
                self.stderr.write(self.style.WARNING(msg))

            except Exception as e:
                msg = f"Erro ao processar {nome_arquivo}: {str(e)}"
                erros.append(msg)
                self.stderr.write(self.style.ERROR(msg))

        self.stdout.write(
            self.style.SUCCESS(f"{len(atualizados)} arquivos importados com sucesso.")
        )
        self.stdout.write(self.style.WARNING(f"{len(erros)} arquivos com erro."))

        if erros:
            caminho_log = os.path.join(pasta, "erros_importacao.txt")
            with open(caminho_log, "w", encoding="utf-8") as erro_file:
                for linha in erros:
                    erro_file.write(f"{linha}\n")
            self.stdout.write(
                self.style.WARNING(f"Erros registrados em: {caminho_log}")
            )

        post_save.connect(update_document, sender=DiarioOficial)
        post_delete.connect(delete_document, sender=DiarioOficial)