import os
import imghdr
import shutil
import hashlib
import zipfile
import exifread
import gradio as gr
import pandas as pd
from PIL import Image

TMP_DIR = "./__pycache__"


def get_exif_data(origin_file_path):
    with open(origin_file_path, "rb") as image_file:
        tags = exifread.process_file(image_file)

    output = ""
    for key in tags.keys():
        value = str(tags[key])
        output += "{0}:{1}\n".format(key, value)

    return output


def clear_exif_data(image_path: str, img_mode=None, outdir=""):
    save_path = f"{TMP_DIR}/{outdir}output." + image_path.split(".")[-1]
    try:
        img = Image.open(image_path)
        data = list(img.getdata())
        if img_mode:
            save_path = (
                f"{TMP_DIR}/{outdir}{hashlib.md5(image_path.encode()).hexdigest()}.jpg"
            )
        else:
            img_mode = img.mode

        img_without_exif = Image.new(img_mode, img.size)
        img_without_exif.putdata(data)
        img_without_exif.save(save_path)

    except Exception as e:
        print(f"\n{image_path} Error: {e}")

    return save_path


def unzip_file(zip_path: str, extract_to=f"{TMP_DIR}/inputs"):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)

    # 打开ZIP文件
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        # 解压文件
        zip_ref.extractall(extract_to)

    return extract_to


def find_files(directory: str):
    found_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            fpath = os.path.join(root, file).replace("\\", "/")
            if imghdr.what(fpath) != None:
                found_files.append(fpath)

    return found_files


def compress(folder_path=f"{TMP_DIR}/outputs", zip_file_path=f"{TMP_DIR}/outputs.zip"):
    if not os.path.exists(folder_path):
        print(f"Error: Folder '{folder_path}' does not exist.")
        return

    with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(folder_path):
            for file in files:
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, folder_path)
                zipf.write(
                    file_path,
                    arcname=os.path.join(os.path.basename(folder_path), relative_path),
                )

    return zip_file_path


def infer(image_path: str, original_ext: bool):
    if not image_path or imghdr.what(image_path) == None:
        return None, "Please input a picture!"

    if os.path.exists(TMP_DIR):
        shutil.rmtree(TMP_DIR)

    os.makedirs(TMP_DIR, exist_ok=True)
    return clear_exif_data(
        image_path, img_mode="RGB" if not original_ext else None
    ), get_exif_data(image_path)


def batch_infer(imgs_zip: str, original_ext: bool):
    if not imgs_zip:
        return None, pd.DataFrame([{"Warning": "Please upload pictures zip!"}])

    if os.path.exists(TMP_DIR):
        shutil.rmtree(TMP_DIR)

    os.makedirs(f"{TMP_DIR}/outputs", exist_ok=True)
    extract_to = unzip_file(imgs_zip)
    imgs = find_files(extract_to)
    mode = "RGB" if not original_ext else None
    exifs = []
    for img in imgs:
        clear_exif_data(img, img_mode=mode, outdir="outputs/")
        exifs.append({"filename": os.path.basename(img), "exif": get_exif_data(img)})

    if not exifs:
        exifs = [{"Warning": "No picture in the zip"}]

    return compress(), pd.DataFrame(exifs)


if __name__ == "__main__":
    with gr.Blocks() as iface:
        with gr.Tab("Process single picture"):
            gr.Interface(
                fn=infer,
                inputs=[
                    gr.File(label="Upload picture", file_types=["image/*"]),
                    gr.Checkbox(
                        label="Export original format",
                        value=False,
                    ),
                ],
                outputs=[
                    gr.Image(
                        label="Download cleaned picture",
                        type="filepath",
                    ),
                    gr.Textbox(label="EXIF", show_copy_button=True),
                ],
                allow_flagging="never",
            )

        with gr.Tab("Batch processor"):
            gr.Interface(
                fn=batch_infer,
                inputs=[
                    gr.File(
                        label="Upload pictures zip",
                        file_types=[".zip"],
                    ),
                    gr.Checkbox(
                        label="Export original format",
                        value=False,
                    ),
                ],
                outputs=[
                    gr.File(
                        label="Download cleaned pictures",
                        type="filepath",
                    ),
                    gr.Dataframe(label="EXIF list"),
                ],
                description="When uploading pictures zip, please make sure the zip is completely uploaded before clicking Submit",
                allow_flagging="never",
            )

    iface.launch()