import argparse
import logging
import multiprocessing
import os
from typing import Iterable

from psd_tools import PSDImage
from tqdm import tqdm

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def parse_args():
    parser = argparse.ArgumentParser(description="Convert PSD files to PNG.")
    parser.add_argument(
        "-d",
        "--directory",
        type=str,
        default="./",
        help="Directory to search for PSD files.",
    )
    parser.add_argument("-o", "--output", type=str, default="./", help="Directory to save PNG files.")
    parser.add_argument(
        "--visible_layers",
        default=[],
        nargs="+",
        type=str,
        help="List of layer names to make visible.",
    )
    parser.add_argument(
        "--invisible_layers",
        default=[],
        nargs="+",
        type=str,
        help="List of layer names to make invisible.",
    )
    parser.add_argument("--num_processes", "-n", default=None, type=int, help=" Number of processes to use.")
    return parser.parse_args()


def find_psd_files(directory):
    psd_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".psd"):
                psd_files.append(os.path.join(root, file))
    return psd_files


def set_layer_visibility(layer, visible_layers, invisible_layers):
    if layer.name in visible_layers:
        layer.visible = True
    if layer.name in invisible_layers:
        layer.visible = False

    if isinstance(layer, Iterable):
        for child in layer:
            set_layer_visibility(child, visible_layers, invisible_layers)


def process_psd_file(task):
    """
    Worker function that processes a single PSD file.
    Opens the PSD, sets layer visibility, composites the image and saves it as PNG.
    """
    psd_file, output, visible_layers, invisible_layers, force = task
    try:
        psd = PSDImage.open(psd_file)
        if force:
            for layer in psd:
                set_layer_visibility(layer, visible_layers, invisible_layers)
        image = psd.composite(force=force)
        fname = os.path.basename(psd_file).replace(".psd", ".png")
        output_file = os.path.join(output, fname)
        image.save(output_file)
    except Exception as e:
        logger.error("Error processing file %s: %s", psd_file, e)


def main(args):
    # Create output directory if it doesn't exist
    if not os.path.exists(args.output):
        os.makedirs(args.output)

    psd_files = find_psd_files(args.directory)
    # force=True when any layer visibility is provided
    force = True if len(args.visible_layers) + len(args.invisible_layers) else False

    tasks = [(psd_file, args.output, args.visible_layers, args.invisible_layers, force) for psd_file in psd_files]

    num_processes = args.num_processes if args.num_processes else multiprocessing.cpu_count() // 2
    # Use multiprocessing to process PSD files in parallel
    with multiprocessing.Pool(processes=num_processes) as pool:
        list(
            tqdm(
                pool.imap_unordered(process_psd_file, tasks),
                total=len(tasks),
                desc="Convert PSD to PNG files",
            )
        )


if __name__ == "__main__":
    args = parse_args()
    main(args)