import gradio as gr
import subprocess
import sys


# ============================================================
# EXECUÇÃO DOS SCRIPTS DO PROJETO
# ============================================================

def executar_script(script):
    try:
        processo = subprocess.Popen(
            [sys.executable, script],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )

        saida = ""

        for linha in processo.stdout:
            saida += linha
            yield saida

        processo.wait()

        if processo.returncode != 0:
            yield (
                saida
                + f"\n\n❌ Processo encerrado com código "
                f"{processo.returncode}."
            )
        else:
            yield saida

    except Exception as e:
        yield f"❌ ERRO INESPERADO:\n\n{e}"


def explorar_corpus():
    yield from executar_script("src/01_explore_corpus.py")


def preparar_dataset():
    yield from executar_script("src/02_prepare_dataset.py")


def treinar_modelo():
    yield from executar_script("src/03_train.py")


def treinar_nllb():
    yield from executar_script("src/03_train_nllb.py")


def avaliar_modelo():
    yield from executar_script("src/04_evaluate.py")


# ============================================================
# DIAGNÓSTICO DO NLLB
# ============================================================

def diagnosticar_nllb():
    """Verifica a configuração e a tokenização do NLLB."""

    yield (
        "🔬 Iniciando diagnóstico do NLLB...\n\n"
        "Carregando configuração e tokenizer. "
        "Isso pode levar alguns minutos na primeira execução."
    )

    try:
        from transformers import AutoConfig, AutoTokenizer

        modelo = "facebook/nllb-200-distilled-600M"
        idioma_origem = "eng_Latn"
        idioma_destino = "mbc_Latn"

        config = AutoConfig.from_pretrained(modelo)

        tokenizer = AutoTokenizer.from_pretrained(
            modelo,
            src_lang=idioma_origem,
        )

        linhas = [
            "✅ Configuração e tokenizer carregados!",
            "",
            "=== 1. ARQUITETURA ===",
            f"Modelo: {modelo}",
            f"Arquitetura: {config.architectures}",
            f"Vocabulário declarado: {config.vocab_size:,}",
            f"Camadas encoder: {getattr(config, 'encoder_layers', 'n/d')}",
            f"Camadas decoder: {getattr(config, 'decoder_layers', 'n/d')}",
            f"Dimensão do modelo: {getattr(config, 'd_model', 'n/d')}",
            "",
            "=== 2. IDIOMA MACUXI ===",
        ]

        mapa_idiomas = getattr(tokenizer, "lang_code_to_id", {})
        idioma_registrado = idioma_destino in mapa_idiomas

        linhas.append(
            f"Código {idioma_destino} registrado: {idioma_registrado}"
        )

        if idioma_registrado:
            linhas.append(
                "O tokenizer já registra esse código de idioma."
            )
        else:
            linhas.extend([
                "Macuxi não está registrado como idioma nativo do NLLB.",
                "Será necessário investigar a adaptação do modelo.",
                "Adicionar apenas um token não ensina o idioma ao modelo.",
            ])

        exemplos = [
            (
                "Always rejoice.",
                "Taatausinpai ako'mantî.",
            ),
            (
                "Brothers, pray for us.",
                "Uyonpayamî', anna pona nîrî epîrematî.",
            ),
        ]

        linhas.extend(["", "=== 3. TOKENIZAÇÃO DOS EXEMPLOS ==="])

        for numero, (ingles, macuxi) in enumerate(exemplos, start=1):
            tokens_en = tokenizer(
                ingles,
                add_special_tokens=True,
            )["input_ids"]

            tokens_mb = tokenizer(
                macuxi,
                add_special_tokens=False,
            )["input_ids"]

            linhas.extend([
                "",
                f"Exemplo {numero}",
                f"Inglês: {ingles}",
                f"Macuxi: {macuxi}",
                f"Quantidade de tokens em inglês: {len(tokens_en)}",
                f"Quantidade de tokens em macuxi: {len(tokens_mb)}",
                "Tokens macuxi: "
                + str(tokenizer.convert_ids_to_tokens(tokens_mb[:30])),
            ])

        linhas.extend([
            "",
            "=== 4. CONCLUSÃO PRELIMINAR ===",
            "Este diagnóstico verifica a configuração e a tokenização.",
            "Ele NÃO comprova que o NLLB consegue traduzir para macuxi.",
            "Nenhum treinamento foi executado.",
            "Os pesos completos do modelo não foram solicitados por este código.",
        ])

        yield "\n".join(linhas)

    except Exception as e:
        yield (
            "❌ ERRO DURANTE O DIAGNÓSTICO\n\n"
            f"Tipo: {type(e).__name__}\n"
            f"Detalhes: {e}\n\n"
            "Verifique as dependências e os logs do Space."
        )


# ============================================================
# TESTE DE COMPATIBILIDADE DO NLLB
# Esta função precisa estar definida ANTES da interface Gradio.
# ============================================================

def testar_compatibilidade_nllb():
    import gc

    modelo_nome = "facebook/nllb-200-distilled-600M"

    try:
        import torch
        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

        yield (
            "🔬 TESTE DE COMPATIBILIDADE NLLB\n\n"
            "Etapa 1/4: carregando tokenizer..."
        )

        tokenizer = AutoTokenizer.from_pretrained(
            modelo_nome,
            src_lang="eng_Latn",
        )

        yield (
            "Etapa 2/4: baixando e carregando os pesos completos.\n"
            "O modelo tem aproximadamente 600 milhões de parâmetros.\n"
            "Esta etapa pode demorar e consumir vários GB de memória."
        )

        modelo = AutoModelForSeq2SeqLM.from_pretrained(
            modelo_nome,
            torch_dtype=torch.float32,
        )

        dispositivo = torch.device("cpu")
        modelo.to(dispositivo)

        yield (
            "Pesos carregados com sucesso!\n\n"
            "Etapa 3/4: adicionando o token de macuxi..."
        )

        token_macuxi = "mbc_Latn"

        quantidade_adicionada = tokenizer.add_special_tokens({
            "additional_special_tokens": [token_macuxi]
        })

        if quantidade_adicionada:
            modelo.resize_token_embeddings(len(tokenizer))

        id_macuxi = tokenizer.convert_tokens_to_ids(token_macuxi)

        if id_macuxi == tokenizer.unk_token_id:
            raise RuntimeError(
                "O tokenizer não conseguiu registrar o token de macuxi."
            )

        modelo.config.vocab_size = len(tokenizer)

        yield (
            "Token de macuxi registrado.\n\n"
            "Etapa 4/4: executando uma passagem de treinamento "
            "com um único exemplo..."
        )

        ingles = "Always rejoice."
        macuxi = "Taatausinpai ako'mantî."

        entradas = tokenizer(
            ingles,
            return_tensors="pt",
            truncation=True,
            max_length=48,
        )

        entradas = {
            chave: valor.to(dispositivo)
            for chave, valor in entradas.items()
        }

        tokens_destino = tokenizer(
            macuxi,
            add_special_tokens=False,
            truncation=True,
            max_length=48,
        )["input_ids"]

        # Convenção de destino: token do idioma + texto + EOS.
        rotulos = torch.tensor(
            [[id_macuxi] + tokens_destino + [tokenizer.eos_token_id]],
            dtype=torch.long,
            device=dispositivo,
        )

        modelo.train()
        modelo.zero_grad(set_to_none=True)

        saida = modelo(**entradas, labels=rotulos)
        perda = saida.loss

        if not torch.isfinite(perda).item():
            raise RuntimeError(
                f"A loss não é finita: {perda.item()}"
            )

        perda.backward()

        gradientes_ok = any(
            parametro.grad is not None
            and torch.isfinite(parametro.grad).all().item()
            for parametro in modelo.parameters()
            if parametro.requires_grad
        )

        modelo.zero_grad(set_to_none=True)

        if not gradientes_ok:
            raise RuntimeError(
                "O modelo não produziu gradientes válidos."
            )

        yield (
            "✅ TESTE CONCLUÍDO COM SUCESSO\n\n"
            f"Modelo: {modelo_nome}\n"
            f"Dispositivo: {dispositivo}\n"
            f"Token de macuxi: {token_macuxi}\n"
            f"ID do token: {id_macuxi}\n"
            f"Loss do exemplo: {perda.item():.4f}\n"
            f"Gradientes válidos: {gradientes_ok}\n\n"
            "Interpretação:\n"
            "- Os pesos foram carregados.\n"
            "- O vocabulário foi ampliado quando necessário.\n"
            "- O modelo executou forward e backward.\n"
            "- Nenhuma atualização de pesos foi realizada.\n\n"
            "ATENÇÃO: isso não comprova capacidade de tradução. "
            "O token novo ainda não aprendeu a representar macuxi, "
            "e será necessário planejar o treinamento experimental."
        )

        del saida, modelo, tokenizer
        gc.collect()

    except Exception as erro:
        gc.collect()

        yield (
            "❌ FALHA NO TESTE DE COMPATIBILIDADE\n\n"
            f"Tipo: {type(erro).__name__}\n"
            f"Detalhes: {erro}\n\n"
            "Nenhum modelo treinado foi salvo. "
            "Envie esta mensagem para investigarmos."
        )


# ============================================================
# INTERFACE GRADIO
# ============================================================

with gr.Blocks(title="SEEDS Lab") as demo:

    gr.Markdown("""
    # 🌱 SEEDS Lab

    ## Pipeline de tradução English → Macushi

    Este espaço executa as etapas do projeto SEEDS.
    """)

    gr.Markdown("### 🔬 Etapas do projeto")

    with gr.Row():
        botao_explorar = gr.Button("🔎 Explorar corpus")
        botao_preparar = gr.Button("📦 Preparar dataset")

    with gr.Row():
        botao_treinar = gr.Button(
            "🧠 Treinar modelo",
            variant="primary",
        )
        botao_avaliar = gr.Button("📊 Avaliar modelo")

    with gr.Row():
        botao_diagnostico = gr.Button("🔬 Diagnosticar NLLB")

    with gr.Row():
        botao_teste_nllb = gr.Button(
            "🧪 Testar compatibilidade NLLB"
        )

    with gr.Row():
        botao_treinar_nllb = gr.Button(
           "🧠 Treinar NLLB + LoRA",
            variant="primary",
        )

    resultado = gr.Textbox(
        label="Resultado",
        lines=30,
        max_lines=50,
    )

    botao_explorar.click(
        fn=explorar_corpus,
        inputs=[],
        outputs=resultado,
    )

    botao_preparar.click(
        fn=preparar_dataset,
        inputs=[],
        outputs=resultado,
    )

    botao_treinar.click(
        fn=treinar_modelo,
        inputs=[],
        outputs=resultado,
    )

    botao_avaliar.click(
        fn=avaliar_modelo,
        inputs=[],
        outputs=resultado,
    )

    botao_diagnostico.click(
        fn=diagnosticar_nllb,
        inputs=[],
        outputs=resultado,
    )

    botao_teste_nllb.click(
        fn=testar_compatibilidade_nllb,
        inputs=[],
        outputs=resultado,
    )

    botao_treinar_nllb.click(
        fn=treinar_nllb,
        inputs=[],
        outputs=resultado,
    )


demo.launch()