main.py

import os
import sys
import fitz  # PyMuPDF
import velopack

from pathlib import Path
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QLabel, QPushButton, QLineEdit, 
                             QFileDialog, QMessageBox, QGroupBox, QRadioButton, 
                             QButtonGroup, QMenu)
from PyQt6.QtCore import Qt
from PyQt6.QtGui import QFont, QAction


class Mode:
    """Represents a PDF processing mode with all associated metadata."""
    
    def __init__(self, name, display_name, section_title, placeholder, help_text, 
                 core_func, parse_input_func, check_overwrite_func,
                 suffix, extension, is_multi_file):
        self.name = name
        self.display_name = display_name
        self.section_title = section_title
        self.placeholder = placeholder
        self.help_text = help_text
        self.core_func = core_func  # The actual PDF manipulation function
        self.parse_input_func = parse_input_func  # Function to parse/validate user input
        self.check_overwrite_func = check_overwrite_func  # Function to check if files exist
        self.suffix = suffix  # Suffix for output files (e.g., "trimmed", "split_part", "img_page")
        self.extension = extension  # File extension (e.g., "pdf", "png")
        self.is_multi_file = is_multi_file  # Whether mode creates multiple files
    
    def get_single_filename(self, base_name):
        """Generate filename for single-file output."""
        return f"{base_name}_{self.suffix}.{self.extension}"
    
    def get_multi_pattern(self, base_name):
        """Generate pattern for multi-file output preview."""
        return f"{base_name}_{self.suffix}*.{self.extension}"
    
    def get_multi_glob_pattern(self, base_name):
        """Generate glob pattern for finding existing multi-file outputs."""
        return f"{base_name}_{self.suffix}*.{self.extension}"


def resource_path(relative_path):
    """ Get absolute path to resource, works for dev and for PyInstaller """
    # PyInstaller creates a temp folder and stores path in _MEIPASS
    base_path = getattr(sys, '_MEIPASS', os.path.abspath("."))
    return os.path.join(base_path, relative_path)


def load_stylesheet():
    """Load the QSS stylesheet from a file."""
    try:
        stylesheet_path = resource_path("style.qss")
        with open(stylesheet_path, "r") as f:
            return f.read()
    except FileNotFoundError:
        print("Warning: style.qss not found. Using default styles.")
        return ""


def parse_page_ranges(page_input):
    """Parse page ranges like '1-3,5,6-9,11' into a sorted list of page numbers."""
    pages = set()
    
    for part in page_input.replace(" ", "").split(","):
        if not part:
            continue
            
        if "-" in part:
            start, end = part.split("-", 1)
            start_page, end_page = int(start), int(end)
            if start_page > end_page:
                raise ValueError(f"Invalid range {part} (start > end)")
            pages.update(range(start_page, end_page + 1))
        else:
            pages.add(int(part))
    
    return sorted(pages)


def parse_chunk_size(chunk_input):
    """Parse and validate chunk size input."""
    try:
        chunk_size = int(chunk_input)
        if chunk_size <= 0:
            raise ValueError("Chunk size must be a positive integer.")
        return chunk_size
    except ValueError as e:
        raise ValueError("Chunk size must be a positive integer.") from e


def check_overwrite_single_file(output_folder, base_name, mode, parsed_input=None):
    """Check if a single output file exists.
    
    Args:
        output_folder: Destination folder path
        base_name: Base filename without suffix/extension
        mode: Mode object with suffix and extension
        parsed_input: Parsed input (not used, for signature compatibility)
    """
    if not output_folder or not base_name:
        return False
    
    filename = mode.get_single_filename(base_name)
    full_path = Path(output_folder) / filename
    return full_path.exists()


def check_overwrite_multi_files(output_folder, base_name, mode, parsed_input=None):
    """Check if any multi-file output files would be overwritten.
    
    Args:
        output_folder: Destination folder path
        base_name: Base filename without suffix/extension
        mode: Mode object with suffix and extension
        parsed_input: Parsed input (not used, for signature compatibility)
    """
    if not output_folder or not base_name:
        return False
    
    output_dir = Path(output_folder)
    pattern = mode.get_multi_glob_pattern(base_name)
    matching_files = list(output_dir.glob(pattern))
    return len(matching_files) > 0


def trim_pdf(input_path, page_numbers, output_folder, base_name, mode):
    """Create a new PDF with only the specified pages.
    
    Args:
        input_path: Source PDF path
        page_numbers: List of page numbers to include
        output_folder: Destination folder
        base_name: Base filename without suffix/extension
        mode: Mode object with suffix and extension
    """
    doc = fitz.open(input_path)
    total_pages = len(doc)
    
    valid_pages = []
    invalid_pages = []
    
    for page_num in page_numbers:
        if 1 <= page_num <= total_pages:
            valid_pages.append(page_num)
        else:
            invalid_pages.append(page_num)
    
    if not valid_pages:
        doc.close()
        raise ValueError("No valid pages to include in the output PDF.")
    
    # Create new document with selected pages (convert to 0-based indexing)
    new_doc = fitz.open()
    for page_num in valid_pages:
        new_doc.insert_pdf(doc, from_page=page_num - 1, to_page=page_num - 1)
    
    # Build output path from folder and base name
    filename = mode.get_single_filename(base_name)
    output_path = Path(output_folder) / filename
    
    new_doc.save(str(output_path))
    new_doc.close()
    doc.close()
    
    message = f"Successfully created PDF with {len(valid_pages)} pages"
    if invalid_pages:
        message += f"\n\nSkipped invalid pages: {invalid_pages} (PDF has {total_pages} pages)"
    
    return message


def split_pdf(input_path, chunk_size, output_folder, base_name, mode):
    """Split a PDF into multiple files with specified chunk size.
    
    Args:
        input_path: Source PDF path
        chunk_size: Number of pages per output file
        output_folder: Destination folder
        base_name: Base filename without suffix/extension
        mode: Mode object with suffix and extension
    """
    doc = fitz.open(input_path)
    total_pages = len(doc)

    if chunk_size <= 0:
        doc.close()
        raise ValueError("Chunk size must be a positive integer.")

    output_dir = Path(output_folder)

    created_files = []
    for chunk_num, start_page in enumerate(range(0, total_pages, chunk_size), start=1):
        end_page = min(start_page + chunk_size, total_pages)

        # Create new document for this chunk
        new_doc = fitz.open()
        new_doc.insert_pdf(doc, from_page=start_page, to_page=end_page - 1)

        # Generate output filename using base_name and suffix
        output_filename = output_dir / f"{base_name}_{mode.suffix}{chunk_num}.{mode.extension}"
        new_doc.save(str(output_filename))
        new_doc.close()

        created_files.append(str(output_filename))
    doc.close()

    num_chunks = len(created_files)
    message = f"Successfully split PDF into {num_chunks} file{'s' if num_chunks > 1 else ''}"
    message += f"\n\nCreated {num_chunks} PDF{'s' if num_chunks > 1 else ''} in:\n{output_dir}"

    return message


def convert_to_images(input_path, page_numbers, output_folder, base_name, mode):
    """Convert specified PDF pages to images.
    
    Args:
        input_path: Source PDF path
        page_numbers: List of page numbers to convert
        output_folder: Destination folder
        base_name: Base filename without suffix/extension
        mode: Mode object with suffix and extension
    """
    doc = fitz.open(input_path)
    total_pages = len(doc)
    
    valid_pages = []
    invalid_pages = []
    
    # Validate page numbers
    for page_num in page_numbers:
        if 1 <= page_num <= total_pages:
            valid_pages.append(page_num)
        else:
            invalid_pages.append(page_num)
    
    if not valid_pages:
        doc.close()
        raise ValueError("No valid pages to convert to images.")
    
    output_dir = Path(output_folder)
    
    # Ensure output directory exists
    output_dir.mkdir(parents=True, exist_ok=True)
    
    created_files = []
    
    # Convert pages to images using PyMuPDF
    for page_num in valid_pages:
        # Convert to 0-based index
        page = doc.load_page(page_num - 1)
        
        # Render page to an image (pixmap)
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x scaling for better quality
        
        # Generate output filename using base_name and suffix
        output_filename = output_dir / f"{base_name}_{mode.suffix}{page_num}.{mode.extension}"
        pix.save(str(output_filename))
        created_files.append(str(output_filename))
    
    doc.close()
    
    num_images = len(created_files)
    message = f"Successfully converted {num_images} page{'s' if num_images > 1 else ''} to image{'s' if num_images > 1 else ''}"
    
    if invalid_pages:
        message += f"\n\nSkipped invalid pages: {invalid_pages} (PDF has {total_pages} pages)"
    
    message += f"\n\nCreated {num_images} PNG image{'s' if num_images > 1 else ''} in:\n{output_dir}"
    
    return message


def update_app():
    """Check for updates and apply them if available."""
    try:
        # Check for updates from GitHub releases
        # Point to the latest release download URL
        manager = velopack.UpdateManager("https://github.com/kelltom/ChiselPDF/releases/latest/download/")  # type: ignore[attr-defined]
        update_info = manager.check_for_updates()
        
        if not update_info:
            QMessageBox.information(None, "No Updates", "You're running the latest version!")
            return  # no updates available
        
        # Ask user if they want to update
        reply = QMessageBox.question(
            None, "Update Available",
            f"A new version is available. Would you like to download and install it?\n\nThis will restart the application.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            # Download the updates
            manager.download_updates(update_info)
            
            # Apply the update and restart the app
            manager.apply_updates_and_restart(update_info)
    except Exception as e:
        # Only show error if it's not the "NotInstalled" development mode error
        if "NotInstalled" not in str(e):
            QMessageBox.warning(None, "Update Error", f"Could not check for updates:\n{str(e)}")


class MainApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.input_path = None
        self.output_folder = None  # Store only the destination folder
        
        # Define available modes
        self.modes = [
            Mode(
                name="selection",
                display_name="Selection",
                section_title="Page Ranges",
                placeholder="e.g., 1-3,5,6-9,11",
                help_text="Specify individual pages and/or ranges separated by commas",
                core_func=trim_pdf,
                parse_input_func=parse_page_ranges,
                check_overwrite_func=check_overwrite_single_file,
                suffix="trimmed",
                extension="pdf",
                is_multi_file=False
            ),
            Mode(
                name="split",
                display_name="Split",
                section_title="Chunk Size",
                placeholder="e.g., 10",
                help_text="Number of pages per file (greater than 0)",
                core_func=split_pdf,
                parse_input_func=parse_chunk_size,
                check_overwrite_func=check_overwrite_multi_files,
                suffix="split_part",
                extension="pdf",
                is_multi_file=True
            ),
            Mode(
                name="image",
                display_name="Image",
                section_title="Page Ranges",
                placeholder="e.g., 1-3,5,6-9,11",
                help_text="Specify individual pages and/or ranges separated by commas",
                core_func=convert_to_images,
                parse_input_func=parse_page_ranges,
                check_overwrite_func=check_overwrite_multi_files,
                suffix="img_page",
                extension="png",
                is_multi_file=True
            )
        ]
        
        # Default to first mode in list
        self.current_mode = self.modes[0]
        
        self.init_ui()
        
    def init_ui(self):
        self.setWindowTitle("ChiselPDF Demo")
        self.setMinimumSize(400, 600)
        self.resize(400, 600)
        self.setStyleSheet(load_stylesheet())
        
        # Create menu bar
        self._create_menu_bar()
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        layout.setContentsMargins(20, 20, 20, 20)
        layout.setSpacing(15)
        
        # Title
        title = QLabel("ChiselPDF")
        title.setFont(QFont("", 18, QFont.Weight.Bold))
        title.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(title)
        layout.addSpacing(10)
        
        # Input section
        layout.addWidget(self._create_input_section())
        
        # Mode selection
        layout.addWidget(self._create_mode_section())
        
        # Page ranges section (dynamic based on mode)
        self.page_section = self._create_page_section()
        layout.addWidget(self.page_section)
        
        # Output section
        layout.addWidget(self._create_output_section())
        
        # Process button
        self.process_button = QPushButton("Execute")
        self.process_button.setMinimumHeight(40)
        self.process_button.setFont(QFont("", 11, QFont.Weight.Bold))
        self.process_button.clicked.connect(self.process_pdf)
        self.process_button.setEnabled(False)
        layout.addWidget(self.process_button)
        
        # Status label
        self.status_label = QLabel("")
        self.status_label.setWordWrap(True)
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.status_label)
        
        layout.addStretch()
    
    def _create_menu_bar(self):
        """Create the application menu bar."""
        menubar = self.menuBar()
        assert menubar is not None  # menuBar() always returns QMenuBar in QMainWindow
        
        # Help menu
        help_menu = menubar.addMenu("Help")
        assert help_menu is not None  # addMenu() always returns QMenu when called with string
        
        # Check for Updates action
        update_action = QAction("Check for Updates", self)
        update_action.triggered.connect(update_app)
        help_menu.addAction(update_action)
    
    def _create_input_section(self):
        group = QGroupBox("Input PDF")
        layout = QVBoxLayout()
        
        row = QHBoxLayout()
        self.input_label = QLabel("No file selected")
        self.input_label.setProperty("labelState", "inactive")
        self.input_label.setWordWrap(True)
        row.addWidget(self.input_label, 1)
        
        browse_btn = QPushButton("Browse...")
        browse_btn.setFixedWidth(100)
        browse_btn.clicked.connect(self.browse_input)
        row.addWidget(browse_btn)
        
        layout.addLayout(row)
        
        self.page_info_label = QLabel("")
        self.page_info_label.setObjectName("pageInfoLabel")
        layout.addWidget(self.page_info_label)
        
        group.setLayout(layout)
        return group
    
    def _create_mode_section(self):
        group = QGroupBox("Mode")
        layout = QHBoxLayout()
        
        self.mode_group = QButtonGroup(self)
        
        # Dynamically create radio buttons for each mode
        for i, mode in enumerate(self.modes):
            radio = QRadioButton(mode.display_name)
            if i == 0:  # Select first mode by default
                radio.setChecked(True)
            radio.toggled.connect(lambda checked, m=mode: self._on_mode_changed(checked, m))
            self.mode_group.addButton(radio)
            layout.addWidget(radio)
        
        layout.addStretch()
        
        group.setLayout(layout)
        return group
    
    def _create_page_section(self):
        group = QGroupBox(self.current_mode.section_title)
        layout = QVBoxLayout()
        
        self.page_entry = QLineEdit()
        self.page_entry.setPlaceholderText(self.current_mode.placeholder)
        layout.addWidget(self.page_entry)
        
        self.help_label = QLabel(self.current_mode.help_text)
        self.help_label.setObjectName("helpLabel")
        layout.addWidget(self.help_label)
        
        group.setLayout(layout)
        return group
    
    def _update_page_section_for_mode(self):
        """Update the page section UI based on current mode."""
        self.page_section.setTitle(self.current_mode.section_title)
        self.page_entry.setPlaceholderText(self.current_mode.placeholder)
        self.help_label.setText(self.current_mode.help_text)
        self.page_entry.clear()
    
    def _create_output_section(self):
        group = QGroupBox("Output")
        layout = QVBoxLayout()
        
        row = QHBoxLayout()
        self.output_label = QLabel("No output location selected")
        self.output_label.setProperty("labelState", "inactive")
        self.output_label.setWordWrap(True)
        row.addWidget(self.output_label, 1)
        
        browse_btn = QPushButton("Browse...")
        browse_btn.setFixedWidth(100)
        browse_btn.clicked.connect(self.browse_output)
        row.addWidget(browse_btn)
        
        layout.addLayout(row)
        group.setLayout(layout)
        return group
    
    def _update_label(self, label, text, active=False):
        """Update a label with truncated text and appropriate styling."""
        display_text = text if len(text) < 80 else f"...{text[-77:]}"
        label.setText(display_text)
        # Use dynamic property instead of inline styles
        label.setProperty("labelState", "active" if active else "inactive")
        # Force style refresh
        label.style().unpolish(label)
        label.style().polish(label)
    
    def _set_label_state(self, label, state):
        """Set a label's state property and refresh styling."""
        label.setProperty("labelState", state)
        if style := label.style():
            style.unpolish(label)
            style.polish(label)
    
    def _get_base_name(self):
        """Get base filename from input path."""
        return Path(self.input_path).stem if self.input_path else "output"
    
    def _get_preview_path(self):
        """Generate preview path for UI display."""
        if not self.output_folder:
            return None
        
        base_name = self._get_base_name()
        if self.current_mode.is_multi_file:
            filename = self.current_mode.get_multi_pattern(base_name)
        else:
            filename = self.current_mode.get_single_filename(base_name)
        
        return str(Path(self.output_folder) / filename)
    
    def _on_mode_changed(self, checked, mode):
        """Handle mode radio button changes."""
        if checked:
            self.current_mode = mode
            self._update_page_section_for_mode()
            self._update_output_suggestion()
    
    def _update_output_suggestion(self):
        """Update the output folder suggestion based on input."""
        if not self.input_path:
            return
        
        # Set output folder to same directory as input
        input_file = Path(self.input_path)
        self.output_folder = str(input_file.parent)
        
        self._update_output_preview()
    
    def _update_output_preview(self):
        """Update the output label to show preview."""
        if preview_path := self._get_preview_path():
            self._update_label(self.output_label, preview_path, active=True)
        else:
            self._update_label(self.output_label, "No output location selected")
    
    def _get_pdf_page_count(self, filename):
        """Get the total number of pages in a PDF file."""
        doc = fitz.open(filename)
        total_pages = len(doc)
        doc.close()
        return total_pages
    
    def _load_input_pdf(self, filename):
        """Load input PDF and update UI with file information."""
        total_pages = self._get_pdf_page_count(filename)
        
        self.input_path = filename
        self._update_label(self.input_label, filename, active=True)
        self.page_info_label.setText(f"Total pages: {total_pages}")
        
        # Auto-suggest output path
        self._update_output_suggestion()
        
        self.process_button.setEnabled(True)
    
    def browse_input(self):
        filename, _ = QFileDialog.getOpenFileName(
            self, "Select Input PDF", "", "PDF files (*.pdf);;All files (*.*)"
        )
        
        if not filename:
            return
        
        try:
            self._load_input_pdf(filename)
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Could not read PDF: {str(e)}")
            self.input_path = None
            self._update_label(self.input_label, "No file selected")
            self.page_info_label.setText("")
    
    def browse_output(self):
        if self.input_path:
            input_file = Path(self.input_path)
            initial_dir = str(input_file.parent)
        else:
            initial_dir = str(Path.home())
        
        if folder := QFileDialog.getExistingDirectory(self, "Select Output Folder", initial_dir):
            self.output_folder = folder
            self._update_output_preview()
    
    def process_pdf(self):
        page_input = self.page_entry.text().strip()
        if not page_input:
            # Determine appropriate input type based on current mode
            input_type = self.current_mode.section_title.lower()
            QMessageBox.warning(self, "Error", f"Please enter {input_type}.")
            return

        self.status_label.setText("Processing...")
        self._set_label_state(self.status_label, "processing")
        QApplication.processEvents()

        try:
            # Validate output folder
            self._ensure_output_folder()

            # Parse/validate input using mode-specific parser
            try:
                parsed_input = self.current_mode.parse_input_func(page_input)
            except (ValueError, AttributeError) as e:
                raise ValueError(f"Invalid input:\n{str(e)}") from e

            # Get base name for output files
            base_name = self._get_base_name()

            # Check for file overwrites using mode-specific checker
            if self.current_mode.check_overwrite_func(self.output_folder, base_name, self.current_mode, parsed_input) and not self._confirm_overwrite():
                return

            # Execute the core PDF manipulation function
            message = self.current_mode.core_func(self.input_path, parsed_input, self.output_folder, base_name, self.current_mode)
            self._show_success(message)

        except Exception as e:
            self.status_label.setText("Failed")
            self._set_label_state(self.status_label, "error")
            QMessageBox.critical(self, "Error", str(e))
    
    def _confirm_overwrite(self):
        """Ask user to confirm file overwrite."""
        reply = QMessageBox.question(
            self, "Confirm Overwrite",
            f"Output file(s) already exist.\n\nOverwrite?",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
        )
        return reply == QMessageBox.StandardButton.Yes
    
    def _show_success(self, message):
        """Display success message."""
        self.status_label.setText(f"Success! Saved to: {self.output_folder}")
        self._set_label_state(self.status_label, "success")
        
        # Create custom message box with "Open Folder" button
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle("Success")
        msg_box.setText(message)
        msg_box.setIcon(QMessageBox.Icon.Information)
        
        # Add standard OK button
        msg_box.addButton(QMessageBox.StandardButton.Ok)
        
        # Add custom "Open Folder" button
        open_folder_btn = msg_box.addButton("Open Folder", QMessageBox.ButtonRole.ActionRole)
        
        msg_box.exec()
        
        # Check if user clicked "Open Folder"
        if msg_box.clickedButton() == open_folder_btn:
            self._open_output_folder()
    
    def _open_output_folder(self):
        """Open the output folder in the system file explorer."""
        if not self.output_folder:
            return
        
        import subprocess
        import platform
        
        folder_path = str(Path(self.output_folder).resolve())
        
        try:
            if platform.system() == "Windows":
                os.startfile(folder_path)
            elif platform.system() == "Darwin":  # macOS
                subprocess.run(["open", folder_path], check=True)
            else:  # Linux and others
                subprocess.run(["xdg-open", folder_path], check=True)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Could not open folder:\n{str(e)}")

    def _ensure_output_folder(self):
        """Validate the output folder before writing files."""
        if not self.output_folder:
            raise ValueError("Please choose an output location before processing.")

        if not self.input_path:
            return

        input_path = Path(self.input_path).resolve()
        output_folder = Path(self.output_folder).resolve()
        
        # Check if single-file output would overwrite input
        if not self.current_mode.is_multi_file:
            base_name = self._get_base_name()
            filename = self.current_mode.get_single_filename(base_name)
            potential_output = output_folder / filename
            if input_path == potential_output:
                raise ValueError("Output file must be different from the input file.")


if __name__ == "__main__":
    # Velopack needs to run first - it may quit/restart the process
    velopack.App().run()  # type: ignore[attr-defined]
    
    app = QApplication(sys.argv)
    window = MainApp()
    window.show()
    sys.exit(app.exec())