import os
import sys
from pathlib import Path
from PIL import Image
from multiprocessing import Pool, cpu_count
from tqdm import tqdm
import time

SUPPORTED_INPUT_EXTS = {'.jpg', '.png', '.bmp', '.tiff', '.webp'}
SUPPORTED_OUTPUT_FORMATS = {'jpg', 'png', 'bmp', 'tiff', 'webp'}

def list_images(folder: Path):
    return [f for f in folder.iterdir() if f.suffix.lower() in SUPPORTED_INPUT_EXTS and f.is_file()]

def normalize_format(fmt: str):
    fmt = fmt.lower()
    if fmt == 'jpg':
        return 'jpeg'
    return fmt

def convert_image(args):
    in_path, out_path, out_format = args
    try:
        with Image.open(in_path) as img:
            img.save(out_path, out_format.upper())
        return (in_path.name, True, None)
    except Exception as e:
        return (in_path.name, False, str(e))

def print_usage():
    print("""
Использование:
  python main.py mass <input_folder> <output_format>
  python main.py single <input_file> <output_format>

Поддерживаемые форматы вывода:
  jpg, png, bmp, tiff, webp
""")

def main():
    if len(sys.argv) < 4:
        print_usage()
        return

    mode = sys.argv[1].lower()
    input_path = Path(sys.argv[2]).resolve()
    out_format = normalize_format(sys.argv[3])

    if out_format not in SUPPORTED_OUTPUT_FORMATS:
        print(f"Ошибка: формат '{out_format}' не поддерживается.")
        print_usage()
        return

    output_dir = Path.cwd() / 'output'
    output_dir.mkdir(exist_ok=True)

    start_time = time.perf_counter()

    if mode == 'mass':
        if not input_path.is_dir():
            print(f"Ошибка: {input_path} не папка.")
            return
        imgs = list_images(input_path)
        if not imgs:
            print(f"В папке {input_path} нет поддерживаемых изображений.")
            return
        print(f"Массовая конвертация {len(imgs)} файлов из {input_path} в {output_dir} с форматом {out_format}")

        args_list = []
        for in_path in imgs:
            out_path = output_dir / (in_path.stem + '.' + out_format)
            args_list.append((in_path, out_path, out_format))

        max_workers = cpu_count()

        with Pool(processes=max_workers) as pool:
            results = []
            with tqdm(total=len(args_list), bar_format='{l_bar}{bar} {n_fmt}/{total_fmt}') as pbar:
                for res in pool.imap_unordered(convert_image, args_list):
                    results.append(res)
                    name, success, err = res
                    if success:
                        tqdm.write(f"[OK] {name}")
                    else:
                        tqdm.write(f"[ERR] {name}: {err}")
                    pbar.update()

    elif mode == 'single':
        if not input_path.is_file():
            print(f"Ошибка: {input_path} не файл.")
            return
        out_path = output_dir / (input_path.stem + '.' + out_format)
        print(f"Конвертация файла {input_path} в {out_path}")
        name, success, err = convert_image((input_path, out_path, out_format))
        if success:
            print(f"[OK] {name}")
        else:
            print(f"[ERR] {name}: {err}")

    else:
        print_usage()
        return

    end_time = time.perf_counter()
    elapsed = end_time - start_time
    minutes, seconds = divmod(elapsed, 60)
    milliseconds = (seconds - int(seconds)) * 1000
    print(f"\nКонвертация заняла {int(minutes)} минут {int(seconds)} секунд {int(milliseconds)} миллисекунд.")

if __name__ == '__main__':
    main()