#!pip install -Uqq fastbook
import fastbook
import gradio as gr
fastbook.setup_book()
from fastbook import *
import os
from zipfile import ZipFile
import shutil
import numpy as np
# import re

title="批量下载图片"
description="""
    1. 可以下载不同种类的图片，图片名称使用空格隔开，比如： 张学友 刘德华 黎明 郭富城\n
    2. 下载的数量应该大于0，程序运行的速度与下载的图片数量和种类个数有关\n
    3. 下载的原理是先在网上爬取图片链接，然后根据链接下载。第2步中输入的下载数量限定的是爬取链接的数量\n
       因为有的链接可能是无效链接，所以下载下来的图片数量可能少于输入的
"""
article="<p style='text-align: center'><center><img src='https://visitor-badge.glitch.me/badge?page_id=kingabzpro/Rick_and_Morty_Bot' alt='visitor badge'></center></p>"

def zipDir(dirpath,outfullname):
    zipp=ZipFile(outfullname,'w')
    for path, dirnames,filenames in os.walk(dirpath):
        fpath=path.replace(dirpath,'')
        for filename in filenames:
            zipp.write(os.path.join(path,filename),os.path.join(fpath,filename))
    zipp.close()
    return outfullname

def download_image(name,number,overwrite):
          
    # 因为在使用Huggingface的时候，其不能访问本地电脑，所以直接内置一个地址
    dirs='E:\图片'
#    name=re.sub(r'[，]{1,100}',',',name)
    names=name.split(' ')
    names=[i for i in names if i != '']
    if len(names)==0:
        return f"请输入要下载的图片名称"
    number=int(number)
    if number<1:
        return f"请输入大于0的下载数量"
    if overwrite=='允许重复下载':
        ow=True
    
    if overwrite=='不允许重复下载':
        ow=False
        
    for o in names:   # 依次获取想要下载图片的关键字，这里是名字
        dest=os.path.join(dirs,o)        # 为每个关键字创建一个独立的文件夹
        os.makedirs(dest,exist_ok=True)   # 创建文件夹
        # fastai自带的方法。 {}内是关键字，f'{}'，是一种格式化方式
        # max_images, 是图片的数量
        # 返回图片的链接
        # 使用的搜索引擎是 DuckDuckGo，网址是:https://duckduckgo.com/?q=&atb=v319-6__
        results=search_images_ddg(f'{o}',max_images=number)   
        # fastai自带的方法，使用并行的方式下载
        # dest 下载目录
        # urls 链接
        download_images(dest,urls=results,preserve_filename=ow)
    
    # 统计图片信息    
    files1=os.listdir(dirs)
    files2=[]
    num1=len(files1)
    num2=[]
    for i in range(num1):
        path2=os.path.join(dirs,files1[i])
        files2.append(path2)
        num2.append(len(os.listdir(path2)))
   
    printing=[]
    for i,j in zip(files1,num2):
        pp=f'{i}: 共有{j}张'
        printing.append(pp)
    prt='\n'.join(printing)
    
    num2=np.array(num2)
    
    
    out1=f'一共下载{num2.sum()}张图片\n其中：\n{prt}\n'
    
    out2=zipDir(dirs,'pictures.zip')
    # 删除文件夹
    for i in files2:
        shutil.rmtree(i)
        
    return (out1,out2)
    
    
demo=gr.Interface(
    fn=download_image,
    inputs=[
        gr.Textbox(
            label="要下载的图片名称，比如：刘德华 张学友",
            lines=1,
            placeholder='刘德华 张学友',
            
           
        ),
        gr.Number(value=10, label="下载图片数量，比如：10"),
        gr.Radio(['不允许重复下载','允许重复下载'],label='是否允许重复下载'),
        
    ],

    outputs=[gr.Textbox(
            label="处理结果",
            lines=10,          
        ),
        gr.File(
            label='图片合集压缩包 请点击下载'
        ),
        ],
    title = title,
    description = description,
    article = article
)

demo.launch(enable_queue=True)