#!/usr/bin/env python3
"""
img2cbz - A command-line tool to convert folders of images (including WebP) into CBZ comic archives.
"""

import os
import sys
import re
import argparse
import logging
import shutil
import tempfile
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from PIL import Image

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S"
)
logger = logging.getLogger("img2cbz")

# Natural sort key for sorting filenames alphanumeric style
# e.g., page_1.png, page_2.png, page_10.png
def natural_sort_key(s):
    return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', s)]

def get_image_files(directory, extensions):
    """Scan directory for files matching the allowed extensions and return them sorted naturally."""
    path = Path(directory)
    files = []
    for f in path.iterdir():
        if f.is_file() and f.suffix.lower().lstrip('.') in extensions:
            files.append(f)
    files.sort(key=lambda x: natural_sort_key(x.name))
    return files

def find_image_dirs(root_dir, extensions):
    """Walk directories recursively and return those containing at least one image."""
    image_dirs = []
    for dirpath, _, filenames in os.walk(root_dir):
        has_images = any(
            Path(f).suffix.lower().lstrip('.') in extensions
            for f in filenames
        )
        if has_images:
            image_dirs.append(Path(dirpath).resolve())
    # Sort them to keep order deterministic
    image_dirs.sort(key=lambda p: natural_sort_key(str(p)))
    return image_dirs

def convert_image(src_path, dest_path, target_format, quality=85):
    """Convert an image to a target format, handling transparency and other mode variations."""
    try:
        with Image.open(src_path) as img:
            # Handle alpha channel / transparency for JPEG conversion
            if target_format.lower() in ("jpg", "jpeg") and img.mode in ("RGBA", "LA", "P"):
                if img.mode == "P":
                    img = img.convert("RGBA")
                # Create a solid white background
                bg = Image.new("RGB", img.size, (255, 255, 255))
                if img.mode == "RGBA":
                    bg.paste(img, mask=img.split()[3])
                else:
                    bg.paste(img, mask=img.convert("RGBA").split()[3])
                img = bg
            elif img.mode == "RGBA" and target_format.lower() in ("png", "webp"):
                # PNG and WebP support transparency, keep as is
                pass
            else:
                # Convert grayscale or other modes to RGB for standard format saving
                if img.mode not in ("RGB", "RGBA"):
                    img = img.convert("RGB")
            
            # Save arguments
            save_args = {}
            if target_format.lower() in ("jpg", "jpeg", "webp"):
                save_args["quality"] = quality
                
            # Create target folder if it doesn't exist
            dest_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Map 'JPG' to 'JPEG' for Pillow
            pil_format = target_format.upper()
            if pil_format == "JPG":
                pil_format = "JPEG"
                
            img.save(dest_path, format=pil_format, **save_args)
        return True
    except Exception as e:
        logger.error(f"Failed to convert {src_path.name} to {target_format.upper()}: {e}")
        return False

def process_directory(source_dir, output_dir, extensions, convert_format, quality, overwrite, clean, dry_run, num_jobs):
    """Process a single folder of images and create a CBZ archive."""
    source_path = Path(source_dir).resolve()
    
    # 1. Gather files
    images = get_image_files(source_path, extensions)
    if not images:
        logger.warning(f"No matching images in '{source_path}'")
        return False
    
    # 2. Determine output destination
    cbz_name = f"{source_path.name}.cbz"
    if output_dir:
        dest_cbz = Path(output_dir).resolve() / cbz_name
    else:
        dest_cbz = source_path.parent / cbz_name
        
    logger.info(f"Found {len(images)} images in '{source_path}'")
    logger.info(f"Target CBZ path: '{dest_cbz}'")
    
    if dest_cbz.exists():
        if overwrite:
            if not dry_run:
                logger.info(f"Output file '{dest_cbz}' already exists. Overwriting...")
        else:
            logger.error(f"Output file '{dest_cbz}' already exists. Use --overwrite to replace it.")
            return False
            
    if dry_run:
        logger.info(f"[DRY-RUN] Would create CBZ archive at '{dest_cbz}' with {len(images)} images.")
        if convert_format:
            logger.info(f"[DRY-RUN] Would convert all images to {convert_format.upper()} (quality: {quality}) first.")
        if clean:
            logger.info(f"[DRY-RUN] Would delete source directory '{source_path}' after verification.")
        return True

    # 3. Create ZIP/CBZ
    success = False
    temp_dir = None
    try:
        files_to_archive = []
        if convert_format:
            temp_dir = tempfile.TemporaryDirectory(prefix="img2cbz_")
            logger.info(f"Converting images to {convert_format.upper()} in parallel...")
            
            tasks = []
            for img_path in images:
                # Retain the base filename but use the target format extension
                new_name = f"{img_path.stem}.{convert_format.lower()}"
                dest_path = Path(temp_dir.name) / new_name
                tasks.append((img_path, dest_path, convert_format, quality))
                
            converted_files = []
            failed_count = 0
            max_workers = num_jobs if num_jobs else os.cpu_count() or 4
            
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_img = {
                    executor.submit(convert_image, src, dest, fmt, q): (src, dest)
                    for src, dest, fmt, q in tasks
                }
                
                completed = 0
                total = len(tasks)
                for future in as_completed(future_to_img):
                    src, dest = future_to_img[future]
                    completed += 1
                    try:
                        if future.result():
                            converted_files.append(dest)
                        else:
                            failed_count += 1
                    except Exception as exc:
                        logger.error(f"Image {src.name} generated an exception during conversion: {exc}")
                        failed_count += 1
                        
                    # Periodically report progress
                    if completed % max(1, total // 5) == 0 or completed == total:
                        logger.info(f"  Progress: {completed}/{total} ({completed*100//total}%)")
            
            if failed_count > 0:
                logger.error(f"Failed to convert {failed_count} images. Aborting archive creation for '{source_path}'")
                return False
                
            # Sort converted files naturally to maintain page ordering
            converted_files.sort(key=lambda x: natural_sort_key(x.name))
            files_to_archive = [(f, f.name) for f in converted_files]
        else:
            # Add images directly
            files_to_archive = [(f, f.name) for f in images]
            
        # Write to Zip Archive
        logger.info(f"Writing to archive '{dest_cbz}'...")
        dest_cbz.parent.mkdir(parents=True, exist_ok=True)
        
        # Decide compression
        try:
            import zlib
            compression_type = zipfile.ZIP_DEFLATED
        except ImportError:
            compression_type = zipfile.ZIP_STORED
            logger.warning("zlib module not found. Compression is disabled; creating stored CBZ.")
            
        with zipfile.ZipFile(dest_cbz, 'w', compression_type) as zip_file:
            for file_path, arc_name in files_to_archive:
                zip_file.write(file_path, arcname=arc_name)
                
        # 4. Verify Archive Integrity
        logger.info("Verifying CBZ archive integrity...")
        with zipfile.ZipFile(dest_cbz, 'r') as zip_file:
            bad_file = zip_file.testzip()
            if bad_file:
                raise zipfile.BadZipFile(f"CRC check failed for: {bad_file}")
                
            zip_contents = zip_file.namelist()
            if len(zip_contents) != len(files_to_archive):
                raise zipfile.BadZipFile(
                    f"File count mismatch. Expected {len(files_to_archive)} files, found {len(zip_contents)} in ZIP."
                )
                
        logger.info(f"Successfully created and verified '{dest_cbz}'")
        success = True
        
    except Exception as e:
        logger.error(f"Failed to create CBZ archive for '{source_path}': {e}")
        if dest_cbz.exists():
            try:
                dest_cbz.unlink()
                logger.info(f"Cleaned up partial/corrupted archive '{dest_cbz}'")
            except Exception as cleanup_err:
                logger.error(f"Could not delete partial archive: {cleanup_err}")
        return False
    finally:
        if temp_dir:
            temp_dir.cleanup()
            
    # 5. Clean up source directory if requested and successful
    if success and clean:
        logger.info(f"Cleaning up source directory '{source_path}'...")
        try:
            shutil.rmtree(source_path)
            logger.info(f"Successfully deleted source directory '{source_path}'")
        except Exception as e:
            logger.error(f"Failed to delete source directory '{source_path}': {e}")
            logger.warning("Archive is verified and safe, but source directory could not be cleaned up completely.")
            
    return success

def run_interactive():
    """Prompt the user for inputs interactively when no CLI arguments are supplied."""
    print("==================================================")
    print("        img2cbz: Interactive Archive Creator      ")
    print("==================================================")
    
    # 1. Ask for directory path
    while True:
        d = input("Enter directory path to convert: ").strip()
        if not d:
            print("Error: Directory path cannot be empty.")
            continue
        p = Path(d)
        if not p.exists():
            print(f"Error: Path '{d}' does not exist.")
            continue
        if not p.is_dir():
            print(f"Error: '{d}' is not a directory.")
            continue
        break
        
    dirs = [str(p.resolve())]
    
    # 2. Ask if recursive scan is desired
    recursive = False
    choice = input("Scan subdirectories recursively? (y/N): ").strip().lower()
    if choice in ("y", "yes"):
        recursive = True
        
    # 3. Ask for output directory (optional)
    out_dir = input("Enter output directory (press Enter to save next to source folders): ").strip()
    if out_dir:
        out_dir = str(Path(out_dir).resolve())
    else:
        out_dir = None
        
    # 4. Ask if they want image conversion
    convert = None
    choice = input("Convert images before archiving? (y/N): ").strip().lower()
    if choice in ("y", "yes"):
        while True:
            fmt = input("Choose conversion format (webp, jpg, png) [default: webp]: ").strip().lower()
            if not fmt:
                convert = "webp"
                break
            elif fmt in ("webp", "jpg", "png"):
                convert = fmt
                break
            else:
                print("Error: Invalid choice. Choose webp, jpg, or png.")
    
    # 5. Ask for quality if convert format is webp/jpg
    quality = 85
    if convert in ("webp", "jpg"):
        q_str = input("Enter compression quality (1-100) [default: 85]: ").strip()
        if q_str.isdigit():
            quality = int(q_str)
            
    # 6. Ask for cleanup mode
    clean = False
    choice = input("Delete source folders after successful archiving? (y/N): ").strip().lower()
    if choice in ("y", "yes"):
        # Double check to prevent accidental loss
        double_check = input("Are you absolutely sure you want to delete source directories? (y/N): ").strip().lower()
        if double_check in ("y", "yes"):
            clean = True
            
    # 7. Overwrite existing CBZs?
    overwrite = False
    choice = input("Overwrite existing CBZ archives if they exist? (y/N): ").strip().lower()
    if choice in ("y", "yes"):
        overwrite = True
        
    # Default options for the rest
    extensions = "jpg,jpeg,png,webp,gif,bmp,tiff"
    dry_run = False
    jobs = None
    verbose = False
    quiet = False
    
    print("==================================================")
    
    return argparse.Namespace(
        dirs=dirs,
        recursive=recursive,
        output_dir=out_dir,
        extensions=extensions,
        convert=convert,
        quality=quality,
        overwrite=overwrite,
        clean=clean,
        dry_run=dry_run,
        jobs=jobs,
        verbose=verbose,
        quiet=quiet
    )

def main():
    parser = argparse.ArgumentParser(
        description="Convert folders of image files, including WebP, into CBZ archives."
    )
    parser.add_argument(
        "dirs",
        metavar="DIR",
        nargs="+",
        help="One or more directories containing image files to convert."
    )
    parser.add_argument(
        "-r", "--recursive",
        action="store_true",
        help="Scan subdirectories recursively. Converts each folder that contains images into its own CBZ."
    )
    parser.add_argument(
        "-o", "--output-dir",
        help="Directory to save the generated CBZ files. Defaults to the parent folder of each source directory."
    )
    parser.add_argument(
        "-e", "--extensions",
        default="jpg,jpeg,png,webp,gif,bmp,tiff",
        help="Comma-separated list of image file extensions to search for (default: jpg,jpeg,png,webp,gif,bmp,tiff)."
    )
    parser.add_argument(
        "-c", "--convert",
        choices=["jpg", "png", "webp"],
        help="Convert all images to the specified format before archiving (ideal for resizing or ensuring compatibility)."
    )
    parser.add_argument(
        "-q", "--quality",
        type=int,
        default=85,
        help="Image quality (1-100) if converting to WebP or JPEG (default: 85)."
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing CBZ archives if they already exist."
    )
    parser.add_argument(
        "--clean",
        action="store_true",
        help="Delete the source directory after successful archive creation and verification."
    )
    parser.add_argument(
        "-d", "--dry-run",
        action="store_true",
        help="Perform a dry run. Show what files/folders would be processed without making modifications."
    )
    parser.add_argument(
        "-j", "--jobs",
        type=int,
        default=None,
        help="Number of threads for parallel image conversion (default: CPU core count)."
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Show verbose output (debug level logging)."
    )
    parser.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress all standard logging, showing errors only."
    )
    if len(sys.argv) == 1:
        try:
            args = run_interactive()
        except (KeyboardInterrupt, EOFError):
            print("\nOperation cancelled.")
            sys.exit(0)
    else:
        args = parser.parse_args()
    
    # Configure log levels
    if args.quiet:
        logger.setLevel(logging.ERROR)
    elif args.verbose:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
        
    # Standardize extensions list
    extensions = [ext.strip().lower().lstrip('.') for ext in args.extensions.split(",")]
    
    # Validate directories
    valid_dirs = []
    for d in args.dirs:
        p = Path(d)
        if not p.exists():
            logger.error(f"Path '{d}' does not exist.")
            sys.exit(1)
        if not p.is_dir():
            logger.error(f"Path '{d}' is not a directory.")
            sys.exit(1)
        valid_dirs.append(p)
        
    # Gather target directories to convert
    target_dirs = []
    if args.recursive:
        for p in valid_dirs:
            target_dirs.extend(find_image_dirs(p, extensions))
        # De-duplicate while preserving order
        seen = set()
        target_dirs = [d for d in target_dirs if not (d in seen or seen.add(d))]
    else:
        target_dirs = [p.resolve() for p in valid_dirs]
        
    if not target_dirs:
        logger.warning("No directories to process.")
        sys.exit(0)
        
    logger.info(f"Target directories identified for conversion: {len(target_dirs)}")
    
    # Process each directory
    successful_archives = 0
    failed_archives = 0
    
    for d in target_dirs:
        logger.info("-" * 50)
        try:
            success = process_directory(
                source_dir=d,
                output_dir=args.output_dir,
                extensions=extensions,
                convert_format=args.convert,
                quality=args.quality,
                overwrite=args.overwrite,
                clean=args.clean,
                dry_run=args.dry_run,
                num_jobs=args.jobs
            )
            if success:
                successful_archives += 1
            else:
                failed_archives += 1
        except Exception as e:
            logger.exception(f"Unhandled exception while processing '{d}': {e}")
            failed_archives += 1
            
    logger.info("=" * 50)
    logger.info(f"Summary: {successful_archives} archives successfully created, {failed_archives} failed.")
    
    if failed_archives > 0:
        sys.exit(1)
    sys.exit(0)

if __name__ == "__main__":
    main()
