from huggingface_hub import HfApi
import json
import datetime



# تعریف تسک‌های رایج NLP
nlp_task_list = [
    "text-classification",
    "token-classification",
    "question-answering",
    "summarization",
    "translation",
    "text-generation",
    "fill-mask",
    "zero-shot-classification",
    "feature-extraction",
    "sentence-similarity",
    "text2text-generation",
    "conversational"
]

def persian_model_finder(nlp_tasks,json_file_name):
    
    api = HfApi()
    all_persian_nlp_models_data = []
    seen_model_ids = set() # برای جلوگیری از اضافه شدن مدل‌های تکراری

    print("در حال جستجو و استخراج اطلاعات مدل‌های NLP فارسی...")

    #  فیلتر کردن و پیمایش روی مدل‌ها
    # برای هر تسک NLP، مدل‌های فارسی را جستجو می‌کنیم.
    # محدودیت 500 مدل برای هر تسک در نظر گرفته شده است تا از دانلود بیش از حد جلوگیری شود.
    # اگر می‌خواهید همه مدل‌ها را استخراج کنید، ممکن است نیاز به پیجینیشن (pagination) باشد.
    for task in nlp_tasks:
        print(f"  جستجو برای تسک: {task} (زبان: فارسی)...")
        try:
            models_for_task = api.list_models(
                language="fa",
                task=task,
                sort="downloads",
                direction=-1, # نزولی (از بیشترین دانلود به کمترین)
                limit=None # می‌توانید این عدد را تغییر دهید
            )

            for model_info in models_for_task:
                if model_info.id not in seen_model_ids:
                    try : # اگر از کارت مدل توانست اطلاعات بیشتری به دست بیاورد :
                        model_ = api.model_info(model_info.id) # به دست آوردن شناسه مدل
                        card_data_dict = model_.card_data.to_dict() # از روی کارت مدل که شامل اطلاعات مدل میباشد یک دیکشنری میسازیم
                        model_data = {
                            "model_id": model_info.id,
                            "url": f"https://huggingface.co/{model_info.id}",
                            "downloads": model_info.downloads,
                            "private": model_info.private,
                            "author": model_info.author,
                            "tags": model_info.tags, # شامل زبان‌ها، تسک‌ها، لایبرری‌ها و...
                            "tag_dataset":"-",
                            "tag_base_model":"-",
                            "tag_license":"-",
                            "tag_region":"-",
                            "pipeline_tag": model_info.pipeline_tag, # تسک اصلی مدل که توسط هاب تعیین شده
                            "Likes":model_info.likes,
                            # چهار مورد پایینی از روی دیکشنری کارت مدل خوانده میشود
                            "languages":card_data_dict.get('language', 'N/A'), # زبان هایی که پشتیبانی میشود
                            "library":card_data_dict.get('library', 'N/A'), # کتابخانه های  مورد استفاده
                            "datasets":card_data_dict.get('datasets', 'N/A'), # دیتابیس های مورد استفاده
                            "license":card_data_dict.get('license', 'N/A'),
                            "just_persian" : False
                            
                        }
                        
                        if model_data["library"] == 'N/A': # در بعضی موارد کتابخانه به این نام ('library_name') در دیکشنری کارت مدل ذخیره شده
                            model_data["library"] = card_data_dict.get('library_name', 'N/A')
                        # شرط پایینی ، مواردی که فقط مختص زبان فارسی هستند را در دیکشنری مشخص میکند
                        if len(model_data["languages"]) == 2 and "multilingual" in model_data["languages"] or\
                            len(model_data["languages"]) == 2 and "persian" in model_data["languages"] or\
                            len(model_data["languages"]) == 2 and "farsi" in model_data["languages"] or\
                            len(model_data["languages"]) == 2 and "fas" in model_data["languages"] or\
                            len(model_data["languages"]) == 2 and model_data["languages"]=="fa" or\
                            model_data["languages"] == "persian"  or\
                            model_data["languages"] == "farsi"  or\
                            model_data["languages"] == "fas"  or\
                            model_data["languages"] == "pes"  or\
                            len(model_data["languages"]) == 1 :
                            model_data["just_persian"] = True
                        
                        for value in model_data["tags"]:
                            
                            if "dataset:" in value :
                                if type(model_data["tag_dataset"]) == type(""):
                                    model_data["tag_dataset"] = list(model_data["tag_dataset"])
                                    model_data["tag_dataset"].pop(0)
                                model_data["tag_dataset"].append(f"{str(value).replace("dataset:","")}")
                            
                            if "base_model:" in value :
                                if type(model_data["tag_base_model"]) == type(""):
                                    model_data["tag_base_model"] = list(model_data["tag_base_model"])
                                    model_data["tag_base_model"].pop(0)
                                model_data["tag_base_model"].append(f"{str(value).replace("base_model:","")}")
                            
                            if "region:" in value :
                                model_data["tag_region"]=f"{str(value).replace("region:","")}"
                                
                            if "license:" in value :
                                model_data["tag_license"]=f"{str(value).replace("license:","")}"
                            
                                
                        all_persian_nlp_models_data.append(model_data)
                        seen_model_ids.add(model_info.id)
                        
                    except : # اگر استفاده از کارت مدل با مشکل مواجه شد :
                        model_data = {
                            "model_id": model_info.id,
                            "url": f"https://huggingface.co/{model_info.id}",
                            "downloads": model_info.downloads,
                            "private": model_info.private,
                            "author": model_info.author,
                            "tags": model_info.tags, # شامل زبان‌ها، تسک‌ها، لایبرری‌ها و...
                            "pipeline_tag": model_info.pipeline_tag, # تسک اصلی مدل که توسط هاب تعیین شده
                            "Likes":model_info.likes,
                            "library":model_info.library_name,
                            # افزودن لایسنس اگر موجود باشد
                            "license": model_info.card_data.license if model_info.card_data and model_info.card_data.license else "N/A"
                        }
                        
                        for value in model_data["tags"]:
                            
                            if "dataset:" in value :
                                if type(model_data["tag_dataset"]) == type(""):
                                    model_data["tag_dataset"] = list(model_data["tag_dataset"])
                                    model_data["tag_dataset"].pop(0)
                                model_data["tag_dataset"].append(f"{str(value).replace("dataset:","")}")
                            
                            if "base_model:" in value :
                                if type(model_data["tag_base_model"]) == type(""):
                                    model_data["tag_base_model"] = list(model_data["tag_base_model"])
                                    model_data["tag_base_model"].pop(0)
                                model_data["tag_base_model"].append(f"{str(value).replace("base_model:","")}")
                            
                            if "region:" in value :
                                model_data["tag_region"]=f"{str(value).replace("region:","")}"
                                
                            if "license:" in value :
                                model_data["tag_license"]=f"{str(value).replace("license:","")}"
                                
                        all_persian_nlp_models_data.append(model_data)
                        seen_model_ids.add(model_info.id)
                        
            print(f"    تعداد مدل‌های یافت شده برای تسک '{task}': {len(models_for_task)}")

        except Exception as e:
            print(f"  خطا در جستجو برای تسک {task}: {e}")

    # مرتب‌سازی نهایی مدل‌ها بر اساس تعداد دانلود (کل لیست)
    # این مرحله اطمینان می‌دهد که حتی اگر مدل‌ها از تسک‌های مختلف جمع‌آوری شده باشند،
    # در نهایت بر اساس دانلود مرتب شده باشند.
    all_persian_nlp_models_data_sorted = sorted(all_persian_nlp_models_data, key=lambda x: x['downloads'], reverse=True)

    print(f"\nتعداد کل مدل‌های NLP فارسی منحصربه‌فرد یافت شده: {len(all_persian_nlp_models_data_sorted)}")

    #  ذخیره اطلاعات در یک فایل JSON
    output_json_file = f"{json_file_name}"
    with open(output_json_file, "w", encoding="utf-8") as f:
        json.dump(all_persian_nlp_models_data_sorted, f, ensure_ascii=False, indent=4)

    print(f"اطلاعات {len(all_persian_nlp_models_data_sorted)} مدل در فایل '{output_json_file}' ذخیره شد.")
    
    
    
persian_model_finder(nlp_task_list,"persian_nlp_models_info.json")