#!/usr/bin/env -S uv run

import os
import shutil
import subprocess
from pathlib import Path
from typing import List, Tuple

import click
import cv2
import img2pdf
import numpy as np
from pdf2image import convert_from_path

import settings
from settings import OptimizationLevel
from PIL import Image

class Settings:
    TRIM_PADDING_PIXELS = 20
    MONOCHROME_THRESHOLD = 127
    OPTIPNG_OPTIMIZATION_LEVEL = 7
    PDF_BORDER_SIZE = 50

settings = Settings()

def ensure_temp_dir():
    """Ensure temporary directory exists and return its path."""
    temp_dir = "temp_processed_images"
    os.makedirs(temp_dir, exist_ok=True)
    return temp_dir

def get_temp_files():
    """Get list of temporary PNG files in order."""
    temp_dir = ensure_temp_dir()
    files = [f for f in os.listdir(temp_dir) if f.endswith('.png')]
    files.sort()  # Ensure correct page order
    return [os.path.join(temp_dir, f) for f in files]

@click.group()
def cli():
    """PDF cleaning toolbox for musical scores."""
    pass

@cli.command()
@click.argument('input_pdf', type=click.Path(exists=True))
def extract(input_pdf):
    """Extract pages from PDF to temporary directory."""
    temp_dir = ensure_temp_dir()
    
    # Convert PDF to images
    print(f"Extracting pages from {input_pdf}...")
    pages = convert_from_path(input_pdf, dpi=400)
    
    # Save each page
    for i, page in enumerate(pages):
        # Convert to grayscale
        page = page.convert('L')
        
        output_path = os.path.join(temp_dir, f"page_{i:03d}.png")
        
        # Save initial version
        page.save(output_path, "PNG", optimize=False)
        
        # Trim whitespace
        trim_whitespace(output_path)
        
        # Reload the trimmed image
        page = Image.open(output_path)
        
        # Calculate new height maintaining aspect ratio
        width = 2048
        ratio = width / page.width
        height = int(page.height * ratio)
        
        # Resize using Lanczos
        page = page.resize((width, height), Image.Resampling.LANCZOS)
        
        # Save final version
        page.save(output_path, "PNG", optimize=False)
        print(f"Saved page {i+1}/{len(pages)}")
    
    print(f"Extracted {len(pages)} pages to {temp_dir}/")

def trim_whitespace(image_path: str, is_final: bool = False) -> bool:
    """Remove white space from around the image.
    
    Handles both RGB and RGBA images, treating transparent pixels as white.
    
    Args:
        image_path: Path to the image file
        is_final: Whether this is the final operation on the image
    
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        # Open image with PIL
        image = Image.open(image_path)
        
        # If image has transparency, flatten it first
        if image.mode == 'RGBA':
            # Create a white background
            background = Image.new('RGB', image.size, 'white')
            # Paste using alpha channel as mask
            background.paste(image, mask=image.split()[3])
            image = background
            
        # Convert to grayscale
        image = image.convert('L')
        
        # Threshold to make all light pixels white and everything else black
        # This helps with finding content bounds
        image = image.point(lambda x: 255 if x > 250 else 0)
        
        # Invert so content is white on black background
        image = Image.eval(image, lambda x: 255 - x)
        
        # Get the bounding box of content (now white pixels)
        bbox = image.getbbox()
        if not bbox:
            print(f"Warning: No content found in {image_path}")
            return False
        
        # Add padding
        padding = settings.TRIM_PADDING_PIXELS
        width, height = image.size
        x1, y1, x2, y2 = bbox
        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(width, x2 + padding)
        y2 = min(height, y2 + padding)
        
        # Open original image again and crop it using the bounds
        original = Image.open(image_path)
        cropped = original.crop((x1, y1, x2, y2))
        
        # Save the cropped image, optimizing only if this is the final operation
        cropped.save(image_path, "PNG", optimize=is_final)
        return True
        
    except Exception as e:
        print(f"Error processing {image_path}: {str(e)}")
        return False

@cli.command()
def deskew():
    """Deskew all pages in temporary directory."""
    temp_files = get_temp_files()
    if not temp_files:
        print("No pages found in temporary directory. Run 'extract' first.")
        return
    
    for file_path in temp_files:
        print(f"Deskewing {os.path.basename(file_path)}...")
        
        # Read image
        image = cv2.imread(file_path)
        
        # Convert to grayscale
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Apply threshold to get binary image
        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        
        # Create a rectangular kernel that's wider than it is tall
        # This helps detect horizontal lines
        kernel_length = np.array(binary).shape[1]//80
        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_length, 1))
        
        # Detect horizontal lines
        horizontal_lines = cv2.erode(binary, horizontal_kernel, iterations=3)
        horizontal_lines = cv2.dilate(horizontal_lines, horizontal_kernel, iterations=3)
        
        # Use probabilistic Hough transform to detect line segments
        line_segments = cv2.HoughLinesP(
            cv2.bitwise_not(horizontal_lines),
            rho=1,
            theta=np.pi/180,
            threshold=100,
            minLineLength=binary.shape[1]//4,  # Lines must be at least 1/4 of image width
            maxLineGap=20
        )
        
        if line_segments is not None and len(line_segments) > 0:
            # Calculate angles of detected line segments
            angles = []
            for line in line_segments:
                x1, y1, x2, y2 = line[0]
                if x2 - x1 == 0:  # Avoid division by zero
                    continue
                angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
                # Only consider angles that are close to horizontal
                if abs(angle) < 20:
                    angles.append(angle)
            
            if angles:
                # Use median angle to avoid outliers
                median_angle = np.median(angles)
                
                # Only rotate if the angle is significant but not too large
                if 0.5 < abs(median_angle) < 20:
                    height, width = image.shape[:2]
                    center = (width/2, height/2)
                    rotation_matrix = cv2.getRotationMatrix2D(center, median_angle, 1.0)
                    rotated = cv2.warpAffine(image, rotation_matrix, (width, height),
                                           flags=cv2.INTER_CUBIC,
                                           borderMode=cv2.BORDER_REPLICATE)
                    
                    # Save rotated image
                    cv2.imwrite(file_path, rotated)
                    print(f"  Rotated by {median_angle:.2f} degrees")
                else:
                    print("  No significant rotation needed")
            else:
                print("  No valid horizontal lines found")
        else:
            print("  No line segments detected")

def convert_to_monochrome(image_path: str, is_final: bool = False) -> bool:
    """Convert image to 1-bit monochrome.
    
    Handles RGBA images by converting transparent pixels to white before thresholding.
    
    Args:
        image_path: Path to the image file
        is_final: Whether this is the final operation on the image
    
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        # Open image with PIL
        image = Image.open(image_path)
        
        # Convert to RGBA if not already
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
        
        # Get the image data as a list of pixels
        data = image.getdata()
        
        # Create new image data, replacing transparent pixels with white
        new_data = []
        for item in data:
            # If pixel is transparent (alpha < 128), make it white
            if item[3] < 128:
                new_data.append((255, 255, 255, 255))
            else:
                new_data.append(item)
        
        # Create new image with modified data
        image.putdata(new_data)
        
        # Convert to grayscale
        image = image.convert('L')
        
        # Convert to 1-bit using threshold
        image = image.point(lambda x: 255 if x > settings.MONOCHROME_THRESHOLD else 0, '1')
        
        # Save the monochrome image, optimizing only if this is the final operation
        image.save(image_path, "PNG", optimize=is_final)
        return True
        
    except Exception as e:
        print(f"Error converting to monochrome {image_path}: {str(e)}")
        return False

def check_optipng_installed() -> bool:
    """Check if optipng is installed."""
    try:
        result = subprocess.run(['optipng', '-v'], capture_output=True, text=True)
        return result.returncode == 0
    except FileNotFoundError:
        return False

def run_optipng(file_path: str) -> bool:
    """Run optipng on a file with error handling.
    
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        result = subprocess.run(
            ['optipng', f'-o{settings.OPTIPNG_OPTIMIZATION_LEVEL}', file_path],
            capture_output=True,
            text=True,
            check=True
        )
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error optimizing {file_path}: {e.stderr}")
        return False
    except Exception as e:
        print(f"Unexpected error optimizing {file_path}: {str(e)}")
        return False

@cli.command()
@click.option('--level', type=click.IntRange(1, 3), default=1,
              help='Optimization level: 1=trim, 2=monochrome, 3=full with optipng')
def optimize(level):
    """Optimize images with specified level of processing.
    
    Optimization Levels:
        1: Only trim whitespace
        2: Level 1 + convert to 1-bit monochrome
        3: Level 2 + optipng optimization
    """
    temp_dir = ensure_temp_dir()
    temp_files = get_temp_files()
    if not temp_files:
        print("No pages found in temporary directory. Run 'extract' first.")
        return
    
    opt_level = OptimizationLevel(level)
    
    total_files = len(temp_files)
    successful = {
        'trim': 0,
        'monochrome': 0,
        'optipng': 0
    }
    
    print(f"Processing {total_files} images at optimization level {level}...")
    
    # Step 1: Convert to monochrome if level >= 2
    if int(opt_level) >= int(OptimizationLevel.MONOCHROME):
        print("\nConverting to monochrome...")
        for i, file_path in enumerate(temp_files, 1):
            print(f"[{i}/{total_files}] Converting {os.path.basename(file_path)}...", end='', flush=True)
            # Only optimize if this is the final step (level 2)
            if convert_to_monochrome(file_path, is_final=(opt_level == OptimizationLevel.MONOCHROME)):
                successful['monochrome'] += 1
                print(" ")
            else:
                print(" ")
    
    # Step 2: Always trim whitespace
    print("\nTrimming whitespace from images...")
    for i, file_path in enumerate(temp_files, 1):
        print(f"[{i}/{total_files}] Processing {os.path.basename(file_path)}...", end='', flush=True)
        # Only optimize if this is the final step (level 1)
        if trim_whitespace(file_path, is_final=(opt_level == OptimizationLevel.TRIM)):
            successful['trim'] += 1
            print(" ")
        else:
            print(" ")
    
    # Step 3: Run optipng if level = 3
    if opt_level == OptimizationLevel.FULL:
        if check_optipng_installed():
            print("\nOptimizing PNG files with optipng...")
            for i, file_path in enumerate(temp_files, 1):
                print(f"[{i}/{total_files}] Optimizing {os.path.basename(file_path)}...", end='', flush=True)
                if run_optipng(file_path):
                    successful['optipng'] += 1
                    print(" ")
                else:
                    print(" ")
        else:
            print("\nNote: optipng not found. Skipping PNG optimization.")
    
    # Print summary
    print("\nOptimization complete!")
    print(f"Successfully trimmed: {successful['trim']}/{total_files} images")
    if int(opt_level) >= int(OptimizationLevel.MONOCHROME):
        print(f"Successfully converted to monochrome: {successful['monochrome']}/{total_files} images")
    if opt_level == OptimizationLevel.FULL and check_optipng_installed():
        print(f"Successfully optimized with optipng: {successful['optipng']}/{total_files} images")

@cli.command()
@click.argument('output_pdf', type=click.Path())
def finalize(output_pdf):
    """Combine processed pages into final PDF and clean up."""
    temp_dir = ensure_temp_dir()
    temp_files = get_temp_files()
    if not temp_files:
        print("No pages found in temporary directory. Run 'extract' first.")
        return
    
    print(f"Combining {len(temp_files)} pages into {output_pdf}...")
    
    # A4 size in millimeters
    A4_WIDTH_MM = 210
    A4_HEIGHT_MM = 297
    
    # Convert to PDF with border
    with open(output_pdf, "wb") as f:
        f.write(img2pdf.convert(
            temp_files,
            with_pdfrw=True,
            layout_fun=img2pdf.get_layout_fun(
                pagesize=(img2pdf.mm_to_pt(A4_WIDTH_MM), img2pdf.mm_to_pt(A4_HEIGHT_MM)),
                border=(settings.PDF_BORDER_SIZE,) * 4,  # Same border size for all sides (top, right, bottom, left)
                fit=img2pdf.FitMode.into
            )
        ))
    
    # Clean up temporary files
    print("Cleaning up temporary files...")
    for file_path in temp_files:
        os.remove(file_path)
    os.rmdir(temp_dir)
    
    print("PDF created successfully and temporary files removed!")

if __name__ == '__main__':
    cli()