import os
import gradio as gr
import shutil
import tempfile
import uuid
from pathlib import Path
import json
from PIL import Image
import fitz  # PyMuPDF for PDF handling

# Constants
TEMP_DIR = "temp"
UPLOAD_DIR = os.path.join(TEMP_DIR, "uploads")
OUTPUT_DIR = os.path.join(TEMP_DIR, "output")
THUMBS_DIR = os.path.join(OUTPUT_DIR, "thumbs")

# Ensure directories exist
for dir_path in [TEMP_DIR, UPLOAD_DIR, OUTPUT_DIR, THUMBS_DIR]:
    os.makedirs(dir_path, exist_ok=True)

def create_thumbnail(image_path, output_path, size=(300, 300)):
    """Create a thumbnail from an image."""
    try:
        with Image.open(image_path) as img:
            img.thumbnail(size, Image.LANCZOS)
            img.save(output_path)
        return output_path
    except Exception as e:
        print(f"Error creating thumbnail: {e}")
        return None

def process_pdf(pdf_path, session_id):
    """Extract pages from a PDF and save as images with thumbnails."""
    pages_info = []
    output_folder = os.path.join(OUTPUT_DIR, session_id)
    thumbs_folder = os.path.join(THUMBS_DIR, session_id)
    
    os.makedirs(output_folder, exist_ok=True)
    os.makedirs(thumbs_folder, exist_ok=True)
    
    try:
        # Open the PDF
        pdf_document = fitz.open(pdf_path)
        
        # Process each page
        for page_num, page in enumerate(pdf_document):
            # Render page to an image with a higher resolution
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
            image_path = os.path.join(output_folder, f"page_{page_num + 1}.png")
            pix.save(image_path)
            
            # Create thumbnail
            thumb_path = os.path.join(thumbs_folder, f"thumb_{page_num + 1}.png")
            create_thumbnail(image_path, thumb_path)
            
            # Add page info
            pages_info.append({
                "src": os.path.join("output", session_id, f"page_{page_num + 1}.png"),
                "thumb": os.path.join("thumbs", session_id, f"thumb_{page_num + 1}.png"),
                "title": f"Page {page_num + 1}"
            })
        
        return pages_info
    except Exception as e:
        print(f"Error processing PDF: {e}")
        return []

def process_images(image_paths, session_id):
    """Process uploaded images and create thumbnails."""
    pages_info = []
    output_folder = os.path.join(OUTPUT_DIR, session_id)
    thumbs_folder = os.path.join(THUMBS_DIR, session_id)
    
    os.makedirs(output_folder, exist_ok=True)
    os.makedirs(thumbs_folder, exist_ok=True)
    
    for i, img_path in enumerate(image_paths):
        try:
            # Copy original image to output folder
            dest_path = os.path.join(output_folder, f"image_{i + 1}.png")
            shutil.copy(img_path, dest_path)
            
            # Create thumbnail
            thumb_path = os.path.join(thumbs_folder, f"thumb_{i + 1}.png")
            create_thumbnail(img_path, thumb_path)
            
            # Add page info
            pages_info.append({
                "src": os.path.join("output", session_id, f"image_{i + 1}.png"),
                "thumb": os.path.join("thumbs", session_id, f"thumb_{i + 1}.png"),
                "title": f"Page {i + 1}"
            })
        except Exception as e:
            print(f"Error processing image {img_path}: {e}")
    
    return pages_info

def create_flipbook(upload_type, pdf_file=None, images=None, view_mode="webgl", skin="light"):
    """Create a flipbook from uploaded PDF or images."""
    try:
        session_id = str(uuid.uuid4())
        pages_info = []
        
        # Process based on upload type
        if upload_type == "pdf" and pdf_file is not None:
            # Save PDF to temp directory
            pdf_path = os.path.join(UPLOAD_DIR, f"{session_id}.pdf")
            with open(pdf_path, "wb") as f:
                f.write(pdf_file)
                
            # Process PDF
            pages_info = process_pdf(pdf_path, session_id)
        elif upload_type == "images" and images is not None:
            # Process images
            pages_info = process_images(images, session_id)
        else:
            return """<div style="color: red; padding: 20px;">Please upload a PDF file or images.</div>"""
        
        if not pages_info:
            return """<div style="color: red; padding: 20px;">Failed to process the uploaded file(s). Please try again.</div>"""
        
        # Create and return HTML for the flipbook
        flipbook_html = generate_flipbook_html(pages_info, session_id, view_mode, skin)
        return flipbook_html
    
    except Exception as e:
        print(f"Error creating flipbook: {e}")
        return f"""<div style="color: red; padding: 20px;">An error occurred: {str(e)}</div>"""

def generate_flipbook_html(pages_info, session_id, view_mode, skin):
    """Generate HTML for the flipbook."""
    # Convert pages_info to JSON for JavaScript
    pages_json = json.dumps(pages_info)
    
    # Create a custom ID for this flipbook
    flipbook_id = f"flipbook_{session_id}"
    
    # HTML template with embedded CSS and JavaScript
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>3D Flipbook</title>
        <style>
            #flipbook-container {{
                width: 100%;
                height: 600px;
                margin: 0 auto;
                position: relative;
            }}
            body, html {{
                margin: 0;
                padding: 0;
                height: 100%;
                overflow: hidden;
            }}
        </style>
        <link rel="stylesheet" type="text/css" href="flipbook.css">
    </head>
    <body>
        <div id="{flipbook_id}" class="flipbook-container"></div>
        
        <script src="flipbook.js"></script>
        <script src="flipbook.webgl.js"></script>
        <script src="flipbook.swipe.js"></script>
        <script src="flipbook.scroll.js"></script>
        <script src="flipbook.book3.js"></script>
        
        <script>
            // Initialize flipbook when page loads
            document.addEventListener('DOMContentLoaded', function() {{
                const options = {{
                    pages: {pages_json},
                    viewMode: '{view_mode}',
                    skin: '{skin}',
                    responsiveView: true,
                    singlePageMode: false,
                    singlePageModeIfMobile: true,
                    pageFlipDuration: 1,
                    sound: false,
                    thumbnailsOnStart: true,
                    btnThumbs: {{ enabled: true }},
                    btnPrint: {{ enabled: false }},
                    btnDownloadPages: {{ enabled: false }},
                    btnDownloadPdf: {{ enabled: false }}
                }};
                
                const container = document.getElementById('{flipbook_id}');
                new FlipBook(container, options);
            }});
        </script>
    </body>
    </html>
    """
    return html

# Define the Gradio interface
with gr.Blocks(title="3D Flipbook Viewer") as demo:
    gr.Markdown("# 3D Flipbook Viewer")
    gr.Markdown("Upload a PDF file or multiple images to create an interactive 3D flipbook.")
    
    with gr.Tabs():
        with gr.TabItem("PDF Upload"):
            pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
            pdf_create_btn = gr.Button("Create Flipbook from PDF")
            
        with gr.TabItem("Image Upload"):
            images = gr.File(label="Upload Images", file_types=["image"], file_count="multiple")
            img_create_btn = gr.Button("Create Flipbook from Images")
    
    with gr.Accordion("Advanced Settings", open=False):
        view_mode = gr.Dropdown(
            choices=["webgl", "3d", "2d", "swipe"], 
            value="webgl", 
            label="View Mode"
        )
        skin = gr.Dropdown(
            choices=["light", "dark", "gradient"], 
            value="light", 
            label="Skin"
        )
    
    output = gr.HTML(label="Flipbook Output")
    
    # Set up event handlers
    pdf_create_btn.click(
        fn=create_flipbook, 
        inputs=[gr.Textbox(value="pdf", visible=False), pdf_file, None, view_mode, skin], 
        outputs=output
    )
    
    img_create_btn.click(
        fn=create_flipbook, 
        inputs=[gr.Textbox(value="images", visible=False), None, images, view_mode, skin], 
        outputs=output
    )
    
# Launch the app
if __name__ == "__main__":
    demo.launch()