import gradio as gr
import pikepdf
import requests
import io
import tempfile
import os

def compress_pdf(input_file, url):
    if input_file is None and (url is None or url.strip() == ""):
        return None, "Please provide either a file or a URL."
    
    if input_file is not None and url and url.strip() != "":
        return None, "Please provide either a file or a URL, not both."
    
    try:
        if url and url.strip() != "":
            response = requests.get(url)
            response.raise_for_status()
            pdf_content = io.BytesIO(response.content)
            initial_size = len(response.content)
        else:
            if hasattr(input_file, 'name'):
                # If input_file is a file path
                pdf_content = input_file.name
                initial_size = os.path.getsize(input_file.name)
            else:
                # If input_file is file-like object
                pdf_content = io.BytesIO(input_file.read())
                pdf_content.seek(0, io.SEEK_END)
                initial_size = pdf_content.tell()
                pdf_content.seek(0)

        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            temp_file_path = temp_file.name

        pdf = pikepdf.Pdf.open(pdf_content)
        
        # Always use the highest compression setting
        compression_params = dict(compress_streams=True, object_stream_mode=pikepdf.ObjectStreamMode.generate)

        pdf.save(temp_file_path, **compression_params)

        # Check the compression ratio achieved
        compressed_size = os.path.getsize(temp_file_path)
        compression_ratio = compressed_size / initial_size
        compression_percentage = (1 - compression_ratio) * 100

        # If compression increased file size or didn't meet minimum threshold, return original file
        if compression_ratio >= 1 or compression_percentage < 5:
            os.remove(temp_file_path)
            return input_file, f"Unable to compress the PDF effectively. Original file returned. (Attempted compression: {compression_percentage:.2f}%)"

        return temp_file_path, f"PDF compressed successfully! Compression achieved: {compression_percentage:.2f}%"
    except Exception as e:
        return None, f"Error compressing PDF: {str(e)}"

def process_and_compress(input_file, url):
    output_file, message = compress_pdf(input_file, url)
    if output_file:
        return output_file, message
    else:
        return None, message

with gr.Blocks() as demo:
    gr.Markdown("# PDF Compressor")
    with gr.Row():
        input_file = gr.File(label="Upload PDF")
        url_input = gr.Textbox(label="Or enter PDF URL")
    compress_btn = gr.Button("Compress")
    output_file = gr.File(label="Download Compressed PDF")
    message = gr.Textbox(label="Message")

    compress_btn.click(
        process_and_compress,
        inputs=[input_file, url_input],
        outputs=[output_file, message]
    )

if __name__ == "__main__":
    demo.launch(share=True)