import wikipedia
import re
import os

def clean_text(text):
    # Elimină referințele și alte elemente nedorite
    text = re.sub(r'\[\d+\]', '', text)
    text = re.sub(r'\n+', '\n', text)
    return text.strip()

def count_words(text):
    return len(text.split())

def create_folder(folder_name):
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)

def save_to_file(folder, filename, content):
    filepath = os.path.join(folder, filename)
    with open(filepath, 'w', encoding='utf-8') as file:
        file.write(content)

def search_wikipedia(keyword, language, num_docs, folder):
    try:
        # Setează limba
        wikipedia.set_lang(language)
        
        # Caută pagini
        search_results = wikipedia.search(keyword, results=num_docs)
        
        if not search_results:
            print(f'No results were found for "{keyword}"')
            return
        
        for i, result in enumerate(search_results, 1):
            try:
                page = wikipedia.page(result)
                content = clean_text(page.content)
                word_count = count_words(content)
                
                # Creează un nume de fișier valid
                filename = f"{i}_{re.sub(r'[^\w\-_\. ]', '_', page.title)}.txt"
                
                # Salvează conținutul în fișier
                save_to_file(folder, filename, content)
                
                print(f"Salvat: {filename} ({word_count} cuvinte)")
            except wikipedia.exceptions.DisambiguationError:
                print(f"Eroare de dezambiguizare pentru: {result}")
            except wikipedia.exceptions.PageError:
                print(f"Pagina nu a fost găsită pentru: {result}")
    except Exception as e:
        print(f"A apărut o eroare: {str(e)}")

def main():
    language = input("Introduceți limba de căutare (ex: ro pentru română, en pentru engleză): ")
    num_docs = int(input("Câte documente doriți să căutați pentru fiecare cuvânt cheie? "))
    folder_name = input("Introduceți numele folderului pentru salvarea rezultatelor: ")
    
    create_folder(folder_name)
    
    while True:
        keyword = input("Introduceți un cuvânt cheie (sau 'q' pentru a ieși): ")
        if keyword.lower() == 'q':
            break
        search_wikipedia(keyword, language, num_docs, folder_name)

if __name__ == "__main__":
    main()