import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from urllib.parse import urlparse
import io

# --- CONFIGURAÇÕES ---
# Seu ID do modelo (Correto conforme você enviou)
MODEL_ID = "asilvamaia/ident_br"

st.set_page_config(page_title="Validador .BR", page_icon="🇧🇷")

# --- FUNÇÃO DE LIMPEZA ---
def limpar_entrada(texto: str) -> str:
    """
    Remove sujeira, portas, protocolos e www.
    """
    texto = str(texto).strip().lower()
    if not texto: return ""
    
    # Mantém e-mail intacto para o modelo rejeitar explicitamente
    if "@" in texto: return texto 

    # Garante protocolo para o urlparse funcionar corretamente
    if "http" not in texto and "://" not in texto:
        texto_temp = "http://" + texto
    else:
        texto_temp = texto

    try:
        parsed = urlparse(texto_temp)
        dominio_limpo = parsed.netloc if parsed.netloc else texto
        
        # Remove a porta (:8080)
        if ":" in dominio_limpo:
            dominio_limpo = dominio_limpo.split(':')[0]
            
        # Remove o 'www.' do início
        if dominio_limpo.startswith("www."):
            dominio_limpo = dominio_limpo[4:]
            
        return dominio_limpo
    except:
        return texto

# --- CARREGAMENTO DO MODELO ---
@st.cache_resource
def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
        # Força uso de CPU no Space gratuito para evitar erros
        model.to("cpu")
        model.eval()
        return tokenizer, model
    except Exception as e:
        st.error(f"Erro ao carregar modelo: {e}")
        return None, None

tokenizer, model = load_model()

# --- INTERFACE DO USUÁRIO ---
st.title("🇧🇷 Validador de Domínios .BR")
st.write("Faça upload de uma lista suja (.txt) para extrair apenas domínios .br válidos.")

uploaded_file = st.file_uploader("Carregar arquivo .txt", type="txt")

if uploaded_file and tokenizer:
    # Lê o arquivo enviado
    stringio = io.StringIO(uploaded_file.getvalue().decode("utf-8"))
    linhas = stringio.readlines()
    
    if st.button(f"Processar {len(linhas)} linhas"):
        validos = []
        rejeitados = []
        
        # Barra de progresso
        progress_bar = st.progress(0)
        total_linhas = len(linhas)
        
        for i, linha in enumerate(linhas):
            original = linha.strip()
            
            # --- CORREÇÃO DA LINHA QUE ESTAVA COM ERRO ---
            # Verifica se a linha é vazia ou se começa com metadados [source: ...]
            if not original or original.startswith("[source:"): 
                continue

            # Aplica limpeza
            limpo = limpar_entrada(original)
            
            # Inferência (Classificação)
            inputs = tokenizer(limpo, return_tensors="pt", truncation=True, max_length=128)
            with torch.no_grad():
                outputs = model(**inputs)
                pred = torch.argmax(outputs.logits, dim=1).item()
            
            # REGRA DE APROVAÇÃO:
            # 1. Modelo diz que é Classe 1 (BR)
            # 2. Tem ponto na string
            # 3. Não tem @
            if pred == 1 and "." in limpo and "@" not in limpo:
                validos.append(limpo)
            else:
                rejeitados.append(original)
            
            # Atualiza barra de progresso a cada 10 itens
            if i % 10 == 0:
                progress_bar.progress(min((i + 1) / total_linhas, 1.0))
        
        progress_bar.progress(1.0)
        
        # --- DEDUPLICAÇÃO (Remove repetidos) ---
        validos_unicos = sorted(list(set(validos)))
        rejeitados_unicos = sorted(list(set(rejeitados)))
        
        st.success("Processamento Concluído!")
        
        # Exibe métricas
        col1, col2 = st.columns(2)
        with col1:
            st.metric("✅ Aprovados (Únicos)", len(validos_unicos), help="Domínios .br válidos")
        with col2:
            st.metric("🔴 Rejeitados (Únicos)", len(rejeitados_unicos), help="Itens descartados")
            
        # Botão de Download
        st.download_button(
            label="⬇️ Baixar Lista Limpa (.txt)",
            data="\n".join(validos_unicos),
            file_name="dominios_limpos.txt",
            mime="text/plain"
        )
        
        # Mostra amostra
        with st.expander("Ver amostra dos resultados"):
            st.write(validos_unicos[:50])