import os
import sys
from PIL import Image
import argparse

def convert_to_webp(input_path, output_path, quality=85, max_width=None, max_height=None):
    try:
        img = Image.open(input_path)
        original_size = os.path.getsize(input_path)
        original_dim = img.size
        
        if max_width or max_height:
            width, height = img.size
            if max_width:
                scale_w = max_width / width
            else:
                scale_w = float('inf')
            if max_height:
                scale_h = max_height / height
            else:
                scale_h = float('inf')
            scale = min(scale_w, scale_h)
            if scale < 1:
                new_width = int(width * scale)
                new_height = int(height * scale)
                img = img.resize((new_width, new_height), Image.LANCZOS)
        
        img.save(output_path, 'WEBP', quality=quality, method=6)
        output_size = os.path.getsize(output_path)
        compression_ratio = (1 - output_size / original_size) * 100
        
        return {
            'success': True,
            'input': input_path,
            'output': output_path,
            'original_size': original_size,
            'output_size': output_size,
            'compression_ratio': compression_ratio,
            'original_dim': original_dim,
            'output_dim': img.size
        }
    except Exception as e:
        return {
            'success': False,
            'input': input_path,
            'error': str(e)
        }

def batch_convert(input_dir, output_dir, thumb_dir=None, quality=85, thumb_quality=80, thumb_width=800):
    results = []
    supported_formats = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.gif')
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    if thumb_dir and not os.path.exists(thumb_dir):
        os.makedirs(thumb_dir)
    
    for filename in sorted(os.listdir(input_dir)):
        if not filename.lower().endswith(supported_formats):
            continue
        
        input_path = os.path.join(input_dir, filename)
        name, _ = os.path.splitext(filename)
        main_output = os.path.join(output_dir, f'{name}.webp')
        
        print(f'Processing: {filename}...')
        
        main_result = convert_to_webp(input_path, main_output, quality=quality)
        results.append(main_result)
        
        if main_result['success']:
            print(f'  Main: {main_result["original_dim"]} -> {main_result["output_dim"]}, '
                  f'{main_result["original_size"]/1024:.1f}KB -> {main_result["output_size"]/1024:.1f}KB '
                  f'({main_result["compression_ratio"]:.1f}% saved)')
            
            if thumb_dir:
                thumb_output = os.path.join(thumb_dir, f'{name}.webp')
                thumb_result = convert_to_webp(input_path, thumb_output, quality=thumb_quality, max_width=thumb_width)
                results.append(thumb_result)
                if thumb_result['success']:
                    print(f'  Thumbnail: {thumb_result["output_dim"]}, '
                          f'{thumb_result["output_size"]/1024:.1f}KB')
    
    return results

def main():
    parser = argparse.ArgumentParser(description='Convert images to WebP format with compression')
    parser.add_argument('input', help='Input file or directory')
    parser.add_argument('-o', '--output', help='Output directory (default: same as input)')
    parser.add_argument('-t', '--thumb', help='Thumbnail output directory')
    parser.add_argument('-q', '--quality', type=int, default=85, help='Main image quality (0-100, default: 85)')
    parser.add_argument('-tq', '--thumb-quality', type=int, default=80, help='Thumbnail quality (0-100, default: 80)')
    parser.add_argument('-tw', '--thumb-width', type=int, default=800, help='Thumbnail max width (default: 800)')
    parser.add_argument('-r', '--recursive', action='store_true', help='Process subdirectories recursively')
    
    args = parser.parse_args()
    
    if os.path.isfile(args.input):
        output_path = args.output
        if not output_path:
            name, _ = os.path.splitext(args.input)
            output_path = f'{name}.webp'
        
        result = convert_to_webp(args.input, output_path, quality=args.quality)
        if result['success']:
            print(f'Converted: {result["input"]}')
            print(f'Original: {result["original_dim"]}, {result["original_size"]/1024:.1f}KB')
            print(f'Output: {result["output_dim"]}, {result["output_size"]/1024:.1f}KB')
            print(f'Saved: {result["compression_ratio"]:.1f}%')
        else:
            print(f'Error: {result["error"]}')
    elif os.path.isdir(args.input):
        output_dir = args.output if args.output else args.input
        
        if args.recursive:
            for root, dirs, files in os.walk(args.input):
                rel_path = os.path.relpath(root, args.input)
                out_dir = os.path.join(output_dir, rel_path)
                thumb_dir = os.path.join(args.thumb, rel_path) if args.thumb else None
                batch_convert(root, out_dir, thumb_dir, args.quality, args.thumb_quality, args.thumb_width)
        else:
            batch_convert(args.input, output_dir, args.thumb, args.quality, args.thumb_quality, args.thumb_width)
    else:
        print(f'Error: {args.input} is not a valid file or directory')
        sys.exit(1)

if __name__ == '__main__':
    main()