import gradio as gr
import os
from zipfile import ZipFile
years = ["1985", "1986", "1987", "1988", "1989", "1990", "1991", "1992", "1993", "1994", "1995", "1996", "1997", "1998", "1999", "2000", "2001", "2002", "2003", "2004", "2005", "2006", "2007", "2008", "2009", "2010", "2011", "2012", "2013", "2014", "2015", "2016", "2017", "2018", "2019", "2020", "2021"]
months = ["all", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"]
files_path = os.path.join("/data", "all_judgement")

def zip_csv_files(directory_path, zip_filename):
    csv_files = [file for file in os.listdir(directory_path) if file.endswith('.csv')]
    with ZipFile(zip_filename, 'w') as zipf:
        for csv_file in csv_files:
            file_path = os.path.join(directory_path, csv_file)
            zipf.write(file_path, csv_file)

def get_file(year_index, month_index):
    file_path = os.path.join(files_path, years[year_index])
    name = str(years[year_index])+'年'+ str(months[month_index] + "月")
    text = "File not exit! Change the path plz, sister.Meng: "
    if months[month_index] == "all":
        directory_to_zip = file_path
        zip_file_name = str(years[year_index])+'年'+ "all.zip"
        zip_csv_files(directory_to_zip, zip_file_name)
        name = zip_file_name
        text = "都压缩起来了奥，文件可能嘎嘎大！稍微等会萌子姐: " + name
        return text, zip_file_name
    else:
        file_path = os.path.join(file_path, name+".csv")

        if not os.path.exists(file_path):
            text = text + name
            return text, os.path.join("/data", "all_judgement", "readme.md")
        else:
            text = "马上就好奥萌子姐: " + name
            return text, file_path

output_component = gr.File()
input_component = [
    gr.Dropdown(years, type="index", label="years", value="2021"),
    gr.Dropdown(months, type="index", label="months", value="all")]
demo = gr.Interface(fn=get_file, inputs = input_component, outputs = ["label", output_component], cache_examples=True, title="输入年月查询")

demo.launch(share=True, server_name='0.0.0.0', server_port=8043, show_api=True)