import requests
import xml.etree.ElementTree as ET
import json
import time
import re

# --- CONFIGURAÇÕES ---
# Endpoint DSpace 7 da UFSCar
OAI_URL = "https://repositorio.ufscar.br/server/oai/request"
TARGET_DATASETS = 50 

# 'oai_dc' é o padrão universal do Dublin Core para interoperabilidade
PARAMS = {
    "verb": "ListRecords",
    "metadataPrefix": "oai_dc"
}

HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36'
}

def remove_namespace(tag):
    """Remove namespaces chatos do XML (ex: {http://purl.org/dc...}title -> title)"""
    return re.sub(r'\{.*?\}', '', tag)

def xml_to_dict(element):
    """
    Converte recursivamente toda a árvore XML do registro para JSON.
    Captura campos repetidos (ex: múltiplos autores ou assuntos) automaticamente.
    """
    tag_clean = remove_namespace(element.tag)
    data = {}
    
    # Atributos (se houver)
    if element.attrib:
        data.update({f"@{k}": v for k, v in element.attrib.items()})

    # Texto
    text = element.text.strip() if element.text else None
    if text:
        data["#text"] = text

    # Filhos
    for child in element:
        child_tag = remove_namespace(child.tag)
        child_data = xml_to_dict(child)
        
        # Se a tag se repete (lista), agrupa em array
        if child_tag in data:
            if not isinstance(data[child_tag], list):
                data[child_tag] = [data[child_tag]]
            data[child_tag].append(child_data)
        else:
            data[child_tag] = child_data
            
    # Simplificação
    if len(data) == 1 and "#text" in data:
        return data["#text"]
    
    return data if data else None

def is_dataset(metadata_dict):
    """
    Detetive: Verifica se o registro é um Dataset.
    O DSpace da UFSCar usa o campo 'type'.
    """
    if not metadata_dict: return False
    
    # Procura o campo 'type' (pode ser lista ou string única)
    types = metadata_dict.get('type', [])
    if isinstance(types, str): types = [types]
    
    # Termos que indicam dados
    keywords = ['dataset', 'dados', 'data', 'conjunto de dados', 'banco de dados']
    
    for t in types:
        if isinstance(t, str):
            if any(k in t.lower() for k in keywords):
                return True
        # Às vezes o xml_to_dict retorna dict se tiver atributo xml:lang
        elif isinstance(t, dict) and '#text' in t:
             if any(k in t['#text'].lower() for k in keywords):
                return True
                
    return False

def coletar_ufscar_completo():
    print(f"1. Conectando ao RI UFSCar ({OAI_URL})...")
    
    datasets_coletados = []
    resumption_token = None
    records_checked = 0
    
    # Loop de Paginação (Necessário pois datasets são raros no meio de teses)
    while len(datasets_coletados) < TARGET_DATASETS:
        
        # Ajusta parâmetros para próxima página se tiver token
        current_params = PARAMS.copy()
        if resumption_token:
            current_params = {"verb": "ListRecords", "resumptionToken": resumption_token}
        
        try:
            print(f"   -> Baixando lote de registros... (Já temos {len(datasets_coletados)} datasets)")
            response = requests.get(OAI_URL, params=current_params, headers=HEADERS, timeout=60)
            
            # Se der erro no servidor, espera um pouco e tenta de novo (simples)
            if response.status_code != 200:
                print(f"Erro no servidor: {response.status_code}. Tentando novamente em 5s...")
                time.sleep(5)
                continue

            root = ET.fromstring(response.content)
            
            # Tratamento de erro OAI (ex: fim da lista ou token expirado)
            error = root.find('.//{http://www.openarchives.org/OAI/2.0/}error')
            if error is not None:
                print(f"Fim da coleta ou erro OAI: {error.text}")
                break

            records = root.findall('.//{http://www.openarchives.org/OAI/2.0/}record')
            
            if not records:
                print("Nenhum registro encontrado neste lote.")
                break
                
            for record in records:
                records_checked += 1
                if len(datasets_coletados) >= TARGET_DATASETS: break
                
                header = record.find('.//{http://www.openarchives.org/OAI/2.0/}header')
                if header.get('status') == 'deleted': continue
                
                metadata_node = record.find('.//{http://www.openarchives.org/OAI/2.0/}metadata')
                
                if metadata_node is not None and len(list(metadata_node)) > 0:
                    # Pega o conteúdo (Dublin Core)
                    content = list(metadata_node)[0]
                    full_data = xml_to_dict(content)
                    
                    # FILTRO: Só salva se for Dataset
                    if is_dataset(full_data):
                        id_oai = header.find('.//{http://www.openarchives.org/OAI/2.0/}identifier').text
                        
                        item_final = {
                            "ufscar_id": id_oai,
                            "metadata_completo": full_data
                        }
                        datasets_coletados.append(item_final)
                        print(f"      [ACHEI!] Dataset encontrado: {id_oai}")

            # Busca token para próxima página
            token_node = root.find('.//{http://www.openarchives.org/OAI/2.0/}resumptionToken')
            if token_node is not None and token_node.text:
                resumption_token = token_node.text
                time.sleep(1) # Pausa educada
            else:
                print("Não há mais páginas no repositório.")
                break
                
        except Exception as e:
            print(f"Erro crítico no lote: {e}")
            break

    # Salva
    nome_arquivo = "ufscar_50_full_metadata.json"
    with open(nome_arquivo, "w", encoding="utf-8") as f:
        json.dump(datasets_coletados, f, ensure_ascii=False, indent=2)
        
    print("\n" + "="*50)
    print(f"SUCESSO! Varri {records_checked} registros para encontrar os {len(datasets_coletados)} datasets.")
    print(f"Arquivo salvo: {nome_arquivo}")
    print("="*50)

if __name__ == "__main__":
    coletar_ufscar_completo()