import argparse
import sys
from pathlib import Path
from PIL import Image

from . import load_color_map, recolor_image, extract_unique_colors


def collect_files(dir_arg: str) -> list[Path]:
    path = Path(dir_arg)

    if "*" in dir_arg:
        parent = path.parent
        pattern = path.name

        if not parent.exists() or not parent.is_dir():
            print("directory does not exist", file=sys.stderr)
            sys.exit(1)

        return [
            f for f in parent.iterdir()
            if f.is_file()
            and f.suffix.lower() == ".png"
            and f.match(pattern)
        ]

    if not path.exists() or not path.is_dir():
        print("directory does not exist", file=sys.stderr)
        sys.exit(1)

    return [
        f for f in path.iterdir()
        if f.is_file() and f.suffix.lower() == ".png"
    ]


def process_file(path: Path, color_map):
    img = Image.open(path)
    img = recolor_image(img, color_map)
    img.save(path)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input")
    parser.add_argument("-dir", "--directory")
    parser.add_argument("-o", "--output")
    parser.add_argument("-r", "--replace", action="store_true")
    parser.add_argument("-m", "--map")
    parser.add_argument("-ext", "--extract", action="store_true")

    args = parser.parse_args()

    if not args.input and not args.directory:
        print("either --input or --directory is required", file=sys.stderr)
        sys.exit(1)

    # -------- directory / glob mode --------
    if args.directory:
        files = collect_files(args.directory)

        if args.extract:
            all_colors = set()

            for file in files:
                img = Image.open(file)
                all_colors |= extract_unique_colors(img)

            for c in sorted(all_colors, key=lambda x: (x.RED, x.GREEN, x.BLUE)):
                print(f"{c.RED:3}, {c.GREEN:3}, {c.BLUE:3}")

            sys.exit(0)

        if not args.map:
            print("color map required for directory mode", file=sys.stderr)
            sys.exit(1)

        try:
            color_map = load_color_map(args.map)
        except Exception as e:
            print(str(e), file=sys.stderr)
            sys.exit(1)

        for file in files:
            process_file(file, color_map)

        sys.exit(0)

    # -------- single file mode --------
    input_path = Path(args.input)
    if not input_path.exists():
        print("input file does not exist", file=sys.stderr)
        sys.exit(1)

    if args.extract:
        img = Image.open(input_path)
        colors = extract_unique_colors(img)

        for c in sorted(colors, key=lambda x: (x.RED, x.GREEN, x.BLUE)):
            print(f"{c.RED:3}, {c.GREEN:3}, {c.BLUE:3}")

        sys.exit(0)

    if args.replace:
        output_path = input_path
    else:
        if not args.output:
            print("output file required unless --replace is set", file=sys.stderr)
            sys.exit(1)
        output_path = Path(args.output)

    if not args.map:
        print("color map required unless --extract is set", file=sys.stderr)
        sys.exit(1)

    try:
        color_map = load_color_map(args.map)
    except Exception as e:
        print(str(e), file=sys.stderr)
        sys.exit(1)

    img = Image.open(input_path)
    img = recolor_image(img, color_map)
    img.save(output_path)


if __name__ == "__main__":
    main()
