import fitz  # PyMuPDF
import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter
import io

# Step 1: Convert PDF to high-quality image
def pdf_to_image(pdf_file, dpi=300):  # Increased DPI for better quality
    try:
        pdf_document = fitz.open(pdf_file.name)
        page = pdf_document.load_page(0)  # Load the first page

        # Render the page at high DPI
        zoom = dpi / 72  # Scale factor based on DPI
        mat = fitz.Matrix(zoom, zoom)
        pix = page.get_pixmap(matrix=mat, alpha=False)

        # Convert pixmap to image
        image_bytes = io.BytesIO(pix.tobytes("png"))
        img = Image.open(image_bytes)
        return img
    except Exception as e:
        print(f"Error during PDF conversion: {str(e)}")
        return None

# Helper: Convert inches to pixels
def inches_to_pixels(inches, dpi=300):
    return int(inches * dpi)

# Step 2: Adjust image dimensions, brightness, thickness, and sharpening
def adjust_image(pdf_file, height_in, width_in, brightness, thickness_scale, sharpening_value):
    img = pdf_to_image(pdf_file)
    if img is None:
        return "Unable to convert PDF to image."

    try:
        # Crop the image to focus on the area of interest (inside the thick black borders)
        # You may need to adjust these coordinates based on your specific diagram layout
        left, top, right, bottom = 50, 50, img.width - 50, img.height - 50  # Example crop values
        img_cropped = img.crop((left, top, right, bottom))

        # Convert dimensions from inches to pixels
        height_px = inches_to_pixels(height_in)
        width_px = inches_to_pixels(width_in)

        # Resize with high-quality resampling
        img_resized = img_cropped.resize((width_px, height_px), Image.LANCZOS)

        # Apply anti-aliasing filter for smooth edges
        img_smoothed = img_resized.filter(ImageFilter.SMOOTH)

        # Adjust brightness
        enhancer = ImageEnhance.Brightness(img_smoothed)
        img_bright = enhancer.enhance(brightness)

        # Apply sharpening
        sharpener = img_bright.filter(ImageFilter.SHARPEN)

        # Save the final image as a PDF without losing quality
        output_pdf_path = "processed_output.pdf"
        sharpener.save(output_pdf_path, "PDF", quality=100)

        return output_pdf_path  # Return the path of the saved PDF
    except Exception as e:
        print(f"Error during image adjustment: {str(e)}")
        return "Error adjusting the image."

# Gradio Interface
iface = gr.Interface(
    fn=adjust_image,
    inputs=[
        gr.File(label="Upload PDF"),
        gr.Number(label="Height (inches)", value=5.0),
        gr.Number(label="Width (inches)", value=4.0),
        gr.Slider(minimum=1, maximum=5, step=0.1, value=1.0, label="Brightness"),
        gr.Slider(minimum=1, maximum=2, step=0.1, value=1.0, label="Thickness Scale"),
        gr.Slider(minimum=1, maximum=5, step=0.1, value=1.0, label="Sharpening"),
    ],
    outputs="file",  # Change output type to 'file' to return the processed PDF
    title="Plate Design Creation - Convert & Adjust PDF Page",
    description="Upload a PDF containing a technical diagram or nameplate. Adjust the dimensions, brightness, thickness, and sharpening for optimal clarity."
)

# Launch the interface
iface.launch()