import os
import tempfile
import math
from typing import List, Tuple
import gradio as gr
from PyPDF2 import PdfReader, PdfWriter

# Load custom CSS
try:
    with open(os.path.join(os.path.dirname(__file__), "fancy.css")) as f:
        custom_css = f.read()
except FileNotFoundError:
    custom_css = ""

# Maximum upload size (1 GB)
MAX_SIZE_BYTES = 1 * 1024 * 1024 * 1024

def split_pdf_into_chunks(file, num_chunks: int) -> Tuple[List[str], str]:
    """
    Teilt eine PDF-Datei in X möglichst gleich große Teile auf.
    Gibt eine Liste der temporären Dateipfade und eine Fehlermeldung (falls vorhanden) zurück.
    """
    if file is None:
        return [], "Bitte lade zuerst eine PDF-Datei hoch."
        
    # Validate file size
    file_size = os.path.getsize(file.name)
    if file_size > MAX_SIZE_BYTES:
        return [], f"Datei überschreitet das 1 GB Limit ({file_size / (1024*1024*1024):.2f} GB)."

    if num_chunks < 1:
        return [], "Die Anzahl der Teile muss mindestens 1 betragen."

    try:
        reader = PdfReader(file.name)
        total_pages = len(reader.pages)
        
        if total_pages == 0:
            return [], "Die PDF-Datei enthält keine Seiten."
            
        # Wenn mehr Teile gefordert sind als Seiten da sind, begrenzen wir es auf die Seitenzahl
        chunks = min(num_chunks, total_pages)
        
        # Berechne die Basis-Seitenanzahl pro Teil und den Rest
        pages_per_chunk = total_pages // chunks
        remainder = total_pages % chunks
        
        output_files = []
        current_page = 0
        
        for i in range(chunks):
            writer = PdfWriter()
            
            # Verteile den Rest auf die ersten Chunks
            actual_pages = pages_per_chunk + (1 if i < remainder else 0)
            
            for _ in range(actual_pages):
                writer.add_page(reader.pages[current_page])
                current_page += 1
                
            # Temporäre Datei für diesen Teil erstellen
            # Verwende ein Muster im Namen, damit man die Reihenfolge erkennt
            out_path = tempfile.NamedTemporaryFile(delete=False, suffix=f"_teil_{i+1}.pdf").name
            with open(out_path, "wb") as f_out:
                writer.write(f_out)
            output_files.append(out_path)
            
        return output_files, ""
        
    except Exception as e:
        return [], f"Fehler beim Aufteilen der PDF: {str(e)}"

def merge_pdfs(files):
    # Merge multiple PDFs in upload order
    if not files:
        return None, "Bitte lade mindestens zwei PDF-Dateien zum Zusammenführen hoch."
    
    # Validate total size
    total_size = sum(os.path.getsize(f.name) for f in files)
    if total_size > MAX_SIZE_BYTES:
        return None, f"Gesamtgröße überschreitet das 1 GB Limit ({total_size / (1024*1024*1024):.2f} GB)."
        
    writer = PdfWriter()
    try:
        for f in files:
            reader = PdfReader(f.name)
            for page in reader.pages:
                writer.add_page(page)
    except Exception as e:
        return None, f"Fehler beim Lesen der PDFs: {e}"
        
    out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
    with open(out_path, "wb") as f_out:
        writer.write(f_out)
    return out_path, None

# Build Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.HTML("<h1 id='header'>PDF Splitter & Merger</h1>")
    gr.Markdown(
        "**Schnelle & Lokale Lösung**: Verarbeite deine PDFs direkt im Speicher ohne dauerhafte Speicherung.\n"
        "Nutze den **Split**-Tab, um eine PDF in gleich große Teile zu zerlegen, oder den **Merge**-Tab, um mehrere PDFs zu vereinen.\n\n"
        "Maximale Dateigröße: **1 GB**."
    )
    
    with gr.Tabs():
        # Split Tab
        with gr.TabItem("PDF Aufteilen"):
            with gr.Row(elem_classes="input-row"):
                pdf_input = gr.File(label="PDF zum Aufteilen auswählen", file_types=['.pdf'])
                chunks_input = gr.Number(label="In wie viele Teile aufteilen?", value=4, precision=0)
            with gr.Row(elem_classes="button-row"):
                split_button = gr.Button("PDF aufteilen", variant="primary")
            
            # Wichtig: file_count="multiple" erlaubt die Anzeige von mehreren Download-Links
            output_split = gr.File(label="Herunterladbare PDF-Teile", file_count="multiple")
            error_split = gr.Textbox(label="Fehlermeldung", interactive=False, visible=False)
            
            def run_split(file, chunks):
                if file is None:
                    return None, "Bitte lade eine PDF-Datei hoch.", True
                if chunks is None or chunks < 1:
                    return None, "Bitte gib eine gültige Anzahl an Teilen ein (mindestens 1).", True
                    
                out_paths, error = split_pdf_into_chunks(file, int(chunks))
                if error:
                    return None, error, True
                return out_paths, "", False
                
            split_button.click(
                fn=run_split,
                inputs=[pdf_input, chunks_input],
                outputs=[output_split, error_split, error_split],
                api_name="split_pdf"
            )
            error_split.change(lambda msg: msg != "", inputs=error_split, outputs=error_split)
            
        # Merge Tab
        with gr.TabItem("PDF Zusammenführen"):
            with gr.Row(elem_classes="input-row"):
                merge_inputs = gr.Files(label="PDF-Dateien zum Zusammenführen auswählen", file_types=['.pdf'])
            with gr.Row(elem_classes="button-row"):
                merge_button = gr.Button("PDFs zusammenführen", variant="primary")
            output_merge = gr.File(label="Zusammengeführte PDF herunterladen")
            error_merge = gr.Textbox(label="Fehlermeldung", interactive=False, visible=False)
            
            def run_merge(files):
                if not files or len(files) < 2:
                    return None, "Bitte lade mindestens zwei PDF-Dateien hoch.", True
                out_path, error = merge_pdfs(files)
                if error:
                    return None, error, True
                return out_path, "", False
                
            merge_button.click(
                fn=run_merge,
                inputs=[merge_inputs],
                outputs=[output_merge, error_merge, error_merge],
                api_name="merge_pdfs"
            )
            error_merge.change(lambda msg: msg != "", inputs=error_merge, outputs=error_merge)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))