#!/usr/bin/env python3

import sys
import os
import argparse
from PIL import Image, ImagePalette

def generate_ansi_256():
    palette = []
    palette.extend([
        (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0),
        (0, 0, 128), (128, 0, 128), (0, 128, 128), (192, 192, 192),
        (128, 128, 128), (255, 0, 0), (0, 255, 0), (255, 255, 0),
        (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255),
    ])
    levels = [0, 95, 135, 175, 215, 255]
    for r in levels:
        for g in levels:
            for b in levels:
                palette.append((r, g, b))
    for i in range(24):
        gray = 8 + i * 10
        palette.append((gray, gray, gray))
    return palette[:256]

def create_palette(colors):
    palette = []
    for r, g, b in colors:
        palette.extend([r, g, b])
    while len(palette) < 768:
        palette.extend([0, 0, 0])
    return palette

def pixelate_and_reduce_colors(image, width, block_scale, dither, palette_type):
    aspect = image.size[1] / image.size[0]
    height = int(width * aspect)

    block_size = max(1, int(width * (block_scale / 100.0)))
    block_width = max(1, width // block_size)
    block_height = max(1, height // block_size)

    image = image.resize((block_width, block_height), Image.NEAREST).convert('RGB')

    if palette_type == 'full':
        return image.resize((block_width * block_size, block_height * block_size), Image.NEAREST).convert('RGB')

    if palette_type == 'ansi256':
        colors = generate_ansi_256()
        pal_image = Image.new('P', (1, 1))
        pal_image.putpalette(create_palette(colors))
        dither_mode = {'floyd': Image.FLOYDSTEINBERG, 'ordered': Image.NONE}[dither]
        image = image.quantize(colors=len(colors), palette=pal_image, method=Image.LIBIMAGEQUANT, dither=dither_mode)
    else:
        raise ValueError(f'[err] unknown palette type: {palette_type}')

    return image.resize((block_width * block_size, block_height * block_size), Image.NEAREST).convert('RGB')

def main():
    parser = argparse.ArgumentParser(description='pixelate an image and reduce colors to a selected palette')
    parser.add_argument('input_image', help='input image file (e.g., png, jpg)')
    parser.add_argument('-o', '--output', default=None, help='output png file (default: same filename in "output" directory)')
    parser.add_argument('--width', type=int, default=800, help='output width in pixels (default: 800)')
    parser.add_argument('--block-scale', type=float, default=1.0,
                        help='block size as a percentage of width (0.1-10.0, default: 1.0)')
    parser.add_argument('--dither', default='floyd', choices=['floyd', 'ordered'],
                        help='dithering method (default: floyd)')
    parser.add_argument('--palette', default='ansi256', choices=['ansi256', 'full'],
                        help='color palette: ansi256 (256 colors), full (16m colors, default: ansi256)')

    args = parser.parse_args()

    if not os.path.isfile(args.input_image):
        print(f'[err] "{args.input_image}" does not exist')
        sys.exit(1)

    if not 0.1 <= args.block_scale <= 10.0:
        print('[err] --block-scale must be between 0.1 and 10.0')
        sys.exit(1)

    if args.width < 1:
        print('[err] --width must be at least 1')
        sys.exit(1)

    try:
        image = Image.open(args.input_image)
        if image.format not in ['PNG', 'JPEG', 'BMP', 'GIF']:
            print(f'[wrn] "{image.format}" may not be fully supported; use png or jpeg for best results')
    except Exception as e:
        print(f'[err] failed to open image: {e}')
        sys.exit(1)

    output_dir = 'output'
    try:
        os.makedirs(output_dir, exist_ok=True)
    except Exception as e:
        print(f'[err] failed to create output directory "{output_dir}": {e}')
        sys.exit(1)

    if args.output:
        output_file = args.output
    else:
        output_file = os.path.join(output_dir, os.path.basename(args.input_image))

    try:
        result = pixelate_and_reduce_colors(
            image, args.width, args.block_scale, args.dither, args.palette
        )
    except Exception as e:
        print(f'[err] failed to process image: {e}')
        sys.exit(1)

    try:
        result.save(output_file)
        print(f'[inf] saved pixelated image to {output_file}')
    except Exception as e:
        print(f'[err] failed to save image: {e}')
        sys.exit(1)

if __name__ == '__main__':
    main()