import hashlib
import json
import os
import pathlib
import re
import struct
import subprocess
from argparse import ArgumentParser

import numpy as np
from PIL import Image
from tqdm import tqdm


def remove_auxiliary_data(input_file, output_file):
    with open(input_file, "rb") as f:
        data = f.read()

    output_data = bytearray()
    output_data.extend(data[:8])  # PNG signature

    # Locate and copy critical chunks
    start = 8
    while start < len(data):
        length, chunk_type = struct.unpack(">I4s", data[start : start + 8])
        start += 8
        chunk_data = data[start : start + length]
        start += length
        crc = data[start : start + 4]
        start += 4

        if chunk_type in [b"IHDR", b"PLTE", b"IDAT", b"IEND"]:
            output_data.extend(struct.pack(">I4s", length, chunk_type))
            output_data.extend(chunk_data)
            output_data.extend(crc)

    with open(output_file, "wb") as f:
        f.write(output_data)


def get_file_id(file_path: str):
    file_id = file_path.replace("\\", "/")
    # assets\\resource\\global\\{SERVERNAME}\\resource\\template\\xxx.png-> {SERVERNAME}/xxx
    m = re.search(r"global/(\w+)/resource/template/([^/]+)\.png", file_id)
    if m:
        return f"{m.group(1)}/{m.group(2)}"
    # assets\\resource\\template\\xxx.png-> official/xxx
    m = re.search(r"template/([^/]+)\.png", file_id)
    if m:
        return f"official/{m.group(1)}"
    # assets\\resource\\PATH\\TO\\xxx.png-> resource/PATH/TO/xxx
    m = re.search(r"resource/(.*)\.png", file_id)
    if m:
        return f"resource/{m.group(1)}"
    # docs\\.vuepress\\public\\image\\PATH\\TO\\xxx.png -> docs/PATH/TO/xxx
    m = re.search(r"docs/.vuepress/public/images/(.*)\.png", file_id)
    if m:
        return f"docs/{m.group(1)}"

    return None


def check_png_need_update(file_path: str, perfect_pngs: dict, quiet: bool):
    # convert to absolute path
    file_path = str(pathlib.Path(file_path).resolve())

    if not file_path.endswith(".png"):
        return False

    file_id = get_file_id(file_path)
    if not file_id:
        if not quiet:
            print("unknown file path", file_path)
        return False

    if file_id in perfect_pngs:
        with open(file_path, "rb") as f:
            data = f.read()
            sha256 = hashlib.sha256(data).hexdigest()
            if sha256 == perfect_pngs[file_id]:
                if not quiet:
                    print("skip", file_path)
                return False

    return True


def update_png_with_optipng(file_path: str, perfect_pngs: dict, quiet: bool):
    if check_png_need_update(file_path, perfect_pngs, quiet):
        if not quiet:
            print("updating", file_path)
    else:
        return False

    file_path = str(pathlib.Path(file_path).resolve())
    file_id = get_file_id(file_path)
    sz_before = os.stat(file_path).st_size

    input_file = output_file = file_path
    remove_auxiliary_data(input_file, output_file)
    if quiet:
        os.system(f'optipng -quiet -o7 -zm1-9 -fix "{output_file}"')
    else:
        os.system(f'optipng -o7 -zm1-9 -fix "{output_file}"')

    sz_after = os.stat(file_path).st_size

    if not quiet:
        print(
            f"before: {sz_before} Bytes, after: {sz_after} Bytes, diff: {sz_before - sz_after} Bytes"
        )

    # update file sha256
    with open(file_path, "rb") as f:
        data = f.read()
        sha256 = hashlib.sha256(data).hexdigest()
        perfect_pngs.update({file_id: sha256})

    if not quiet:
        print("updated", file_path)
    update_perfect_png_dict(perfect_pngs)
    return sz_before - sz_after


def update_png_with_oxipng(file_path: str, perfect_pngs: dict, quiet: bool):
    if check_png_need_update(file_path, perfect_pngs, quiet):
        if not quiet:
            print("updating", file_path)
    else:
        return False

    file_path = str(pathlib.Path(file_path).resolve())
    file_id = get_file_id(file_path)
    sz_before = os.stat(file_path).st_size

    with Image.open(file_path) as img:
        img_before = np.array(img.convert("L"))

    # Build oxipng command (list form is safe without shell escaping)
    oxipng_cmd_base = ["oxipng", "-o", "max", "--fast", "-Z", "-s"]
    if quiet:
        oxipng_cmd_base.append("-q")

    # First pass
    subprocess.run(oxipng_cmd_base + [file_path], check=True)
    # Second pass
    subprocess.run(
        ["oxipng", "-o", "2", "-s"] + (["-q"] if quiet else []) + [file_path],
        check=True,
    )

    sz_after = os.stat(file_path).st_size

    if not quiet:
        print(
            f"before: {sz_before} Bytes, after: {sz_after} Bytes, diff: {sz_before - sz_after} Bytes"
        )

    # update file sha256
    with open(file_path, "rb") as f:
        data = f.read()
        sha256 = hashlib.sha256(data).hexdigest()
        perfect_pngs.update({file_id: sha256})

    with Image.open(file_path) as img:
        img_after = np.array(img.convert("L"))
    if not (img_before == img_after).all():
        return -1

    if not quiet:
        print("updated", file_path)
    update_perfect_png_dict(perfect_pngs)
    return sz_before - sz_after


cur_dir = pathlib.Path(__file__).parent.resolve()
perfect_pngs_path = str(cur_dir / "optimize_templates.json")


def update_perfect_png_dict(perfect_pngs: dict):
    with open(perfect_pngs_path, "w") as f:
        json.dump(perfect_pngs, f, indent=4)
        f.write("\n")  # 保持与 prettier 一致的末尾换行


def ArgParser():
    parser = ArgumentParser()
    parser.add_argument(
        "-p",
        "--path",
        dest="path",
        nargs="*",
        help="the paths of png files or folders",
        default=[],
    )
    parser.add_argument("-q", "--quiet", dest="quiet", action="store_true")
    return parser


if __name__ == "__main__":
    args = ArgParser().parse_args()
    paths = args.path
    quiet = args.quiet

    perfect_pngs = {}
    optimized_png_paths = []

    if os.path.exists(perfect_pngs_path):
        with open(perfect_pngs_path, "r") as f:
            perfect_pngs = json.load(f)

    if len(paths) == 0:
        docs_dir = cur_dir.parent.parent / "docs"
        resource_dir = cur_dir.parent.parent / "assets" / "resource"
        paths = [
            str(resource_dir),
            str(docs_dir),
        ]
        print("no path specified, use default paths:", paths)

    files = []
    tqdm_total = 0

    for path in paths:
        if pathlib.Path(path).is_dir():
            for root, _, _files in os.walk(path):
                for f in _files:
                    file = os.path.join(root, f)
                    if check_png_need_update(file, perfect_pngs, quiet):
                        files.append(file)
                        tqdm_total += os.stat(file).st_size ** 2 // 2**20
        elif pathlib.Path(path).is_file():
            files.append(path)
            tqdm_total += os.stat(path).st_size ** 2 // 2**20

    total_diff_sz = 0
    if quiet:
        t = tqdm(files, total=tqdm_total)
    for i, file in enumerate(files):
        cur_file_sz = os.stat(file).st_size
        diff_sz = update_png_with_oxipng(file, perfect_pngs, quiet)
        if diff_sz == -1:
            print("error.")
            exit(-1)
        if diff_sz:
            total_diff_sz += diff_sz
            optimized_png_paths.append(file)

        if total_diff_sz < 1024:
            total_diff_sz_str = f"{total_diff_sz} Bytes"
        elif total_diff_sz < 1048576:
            total_diff_sz_str = f"{(total_diff_sz / 1024):.2f} KiB"
        else:
            total_diff_sz_str = f"{(total_diff_sz / 1048576):.2f} MiB"

        if quiet:
            t.update(cur_file_sz**2 // 2**20)
            t.set_postfix(
                file_counts=f"{i+1}/{len(files)}", reduced_size=total_diff_sz_str
            )
            t.refresh()
        else:
            print(
                f"file counts: {i+1}/{len(files)}, reduced pngs size: {total_diff_sz_str}"
            )
