from blacklist import bad_tags,bad_models,fav_models
models = []
models_plus_tags=[]
tags_plus_models=[]


def list_uniq(l):
    return sorted(set(l), key=l.index)

def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False, bad_models=bad_models):
    from huggingface_hub import HfApi
    api = HfApi()
    default_tags = ["diffusers"]
    if not sort: sort = "last_modified"
    models = []
    models_plus_tags=[]
    try:
        model_infos = api.list_models(author=author, task="text-to-image",
                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
    except Exception as e:
        print(f"Error: Failed to list models.")
        print(e)
        return models
    for model in model_infos:
        
        if not model.private and not model.gated:
            loadable = True
            if not_tag and not_tag in model.tags or not loadable: continue
            if model.id not in bad_models :
                models.append(model.id)
                #models_plus_tags.append([model.id,process_tags(model.tags)])
                models_plus_tags.append([model.id,process_tags(model.cardData.tags)])
            if len(models) == limit: break
    return models , models_plus_tags

def find_warm_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False, bad_models=bad_models):
    from huggingface_hub import HfApi
    api = HfApi()
    default_tags = ["diffusers"]
    if not sort: sort = "last_modified"
    models = []
    models_plus_tags=[]
    try:
        model_infos = api.list_models(author=author, task="text-to-image",
                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit, inference ="warm")
    except Exception as e:
        print(f"Error: Failed to list models.")
        print(e)
        return models
    for model in model_infos:
        
        if not model.private and not model.gated:
            loadable = True
            if not_tag and not_tag in model.tags or not loadable: continue
            if model.id not in bad_models :
                models.append(model.id)
                #models_plus_tags.append([model.id,process_tags(model.tags)])
                models_plus_tags.append([model.id,process_tags(model.cardData.tags)])
            if len(models) == limit: break
    return models , models_plus_tags

def process_tags(tags,bad_tags=bad_tags):
    t1=True
    new_tags=[]
    for tag in tags:
        if tag not in bad_tags:
            if tag == 'en':
                t1=False
            if t1 :
                new_tags.append(tag)
    return new_tags

def clas_tags(models_plus_tags,min):
    tags_plus_models=[]
    listTemp=[]
    listAllModels=['all',0,[]]
    new_tag=True
    for ent in models_plus_tags:
        listAllModels[1]+=1
        listAllModels[2].append(ent[0])
        for tagClas in ent[1]:
            new_tag=True
            for tag in listTemp:
                if tag[0] == tagClas:
                    tag[1]+=1
                    tag[2].append(ent[0])
                    new_tag=False
            if new_tag:
                listTemp.append([tagClas,1,[ent[0]]])
    tags_plus_models.append(listAllModels)
    for t in listTemp:
        if t[1]>=min and t[1]!=len(models_plus_tags):
            tags_plus_models.append(t)
    return tags_plus_models
    
def orga_tag(tags_plus_models):
    output=[]
    while(len(output)<len(tags_plus_models)):
        max=0
        for tag in tags_plus_models:
            if tag[1]>max:
                if tag not in output:
                    max=tag[1]
                    tagMax=tag
        output.append(tagMax)
    return output
def tag_new(models,limit=40):
    output=["NEW",0,[]]
    if limit>len(models):
        limit=len(models)
    for i in range(limit):
        output[1]+=1
        output[2].append(models[i])
    return output
    
def tag_fav(models=models,fav_models=fav_models):
    output=["FAV",0,[]]
    for m in fav_models:
        if m in models:
            output[1]+=1
            output[2].append(m)
    return output

def update_tag(models_plus_tags,list_new_tag):
    for m_new in list_new_tag[2]:
        for m in models_plus_tags:
            if m_new == m[0]:
                m[1].append(list_new_tag[0])
    return models_plus_tags
    
    
from operator import itemgetter

nb_models_to_load=10000
#models = find_model_list("Yntec", [], "", "last_modified", 20)
#models , models_plus_tags = find_model_list("John6666", ["stable-diffusion-xl"], "", "last_modified", nb_models_to_load)
models , models_plus_tags = find_model_list("John6666", ["stable-diffusion-xl"], "not-for-all-audiences", "last_modified", nb_models_to_load)
#tags_plus_models = orga_tag(clas_tags(models_plus_tags,2))
tags_plus_models = orga_tag(clas_tags(sorted(models_plus_tags, key=itemgetter(0)),2))
list_new=tag_new(models,nb_models_to_load)
models_plus_tags=update_tag(models_plus_tags,list_new)
tags_plus_models.insert(1,list_new)
list_fav=tag_fav(models)
if list_fav[1]>0:
    tags_plus_models.insert(1,list_fav)
    models_plus_tags=update_tag(models_plus_tags,list_fav)

#models.extend(find_model_list("John6666", ["stable-diffusion-xl"], "", "last_modified", 200))

#models.extend(find_model_list("John6666", [], "", "last_modified", 20)) # The latest 20 models will be added to the models written above.

# Examples:
#models = ['yodayo-ai/kivotos-xl-2.0', 'yodayo-ai/holodayo-xl-2.1'] # specific models
#models = find_model_list("Yntec", [], "", "last_modified", 20) # Yntec's latest 20 models
#models = find_model_list("Yntec", ["anime"], "", "last_modified", 20) # Yntec's latest 20 models with 'anime' tag
#models = find_model_list("Yntec", [], "anime", "last_modified", 20) # Yntec's latest 20 models without 'anime' tag
#models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
#models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface