import os
from playwright.sync_api import sync_playwright
from urllib.request import urlretrieve

def main():
    # Cerem utilizatorului să introducă cuvintele cheie pentru căutare
    search_terms = input("Introduceți cuvintele cheie pentru căutare (ex. 'rețele neuronale'): ")
    
    # Cerem utilizatorului să introducă numărul de documente pe care dorește să le descarce
    num_documents = int(input("Introduceți numărul de documente pe care doriți să le descărcați: "))

    # Inițializăm Playwright
    with sync_playwright() as pw:
        # Lansăm browser-ul Firefox (vizibil, cu o mică întârziere între acțiuni)
        browser = pw.firefox.launch(headless=False, slow_mo=2000)
        page = browser.new_page()
        
        # Navigăm la pagina de căutare arXiv
        page.goto("http://arxiv.org/search")
        
        # Completăm câmpul de căutare cu termenii introduși
        search_input = page.get_by_placeholder("Search term")
        search_input.fill(search_terms)
        
        # Apăsăm pe al doilea buton "Search"
        search_button = page.get_by_role("button").get_by_text("Search").nth(1)
        search_button.click()
        
        # Așteptăm să se încarce rezultatele
        page.wait_for_load_state("networkidle")
        
        # Creăm folderul "data" dacă nu există
        if not os.path.exists("data"):
            os.mkdir("data")
        
        # Variabilă pentru a ține evidența documentelor descărcate
        downloaded = 0
        
        # Continuăm până descărcăm numărul cerut de documente
        while downloaded < num_documents:
            # Găsim toate link-urile către PDF-uri pe pagina curentă
            links = page.locator("xpath=//a[contains(@href, 'arxiv.org/pdf')]").all()
            
            # Descărcăm documentele de pe pagina curentă
            for link in links:
                if downloaded >= num_documents:
                    break
                url = link.get_attribute("href")
                if url:
                    # Creăm un nume de fișier bazat pe ultimele 5 caractere din URL
                    file_name = f"data/{url[-5:]}.pdf"
                    # Descărcăm fișierul
                    urlretrieve(url, file_name)
                    print(f"Descărcat: {file_name}")
                    downloaded += 1
            
            # Dacă mai avem documente de descărcat, trecem la pagina următoare
            if downloaded < num_documents:
                next_button = page.locator("xpath=//a[contains(text(), 'Next')]")
                if next_button.count() > 0:
                    next_button.click()
                    page.wait_for_load_state("networkidle")
                else:
                    print("Nu mai sunt pagini disponibile.")
                    break
        
        # Închidem browser-ul
        browser.close()

if __name__ == "__main__":
    main()