#!/usr/bin/env python3

import os
import logging
from typing import Optional, Dict, Any, Tuple
from Quartz import PDFKit, CoreGraphics, kCGPDFContextUserPassword
from Foundation import NSURL

class PDFMergeError(Exception):
    """Custom exception for PDF merge errors"""
    pass

class PDFMerger:
    """Handles merging of letterhead and content PDFs"""
    
    def __init__(self, letterhead_path: str):
        """
        Initialize PDFMerger with a letterhead PDF path
        
        Args:
            letterhead_path: Path to the letterhead PDF template
        """
        self.letterhead_path = os.path.expanduser(letterhead_path)
        logging.info(f"Initializing PDFMerger with template: {self.letterhead_path}")

    def merge(self, input_path: str, output_path: str, strategy: str = "darken") -> None:
        """
        Merge letterhead with input PDF using the specified strategy
        
        Args:
            input_path: Path to the content PDF
            output_path: Path to save the merged PDF
            strategy: Merging strategy to use (overlay, multiply, transparency, etc.)
        
        Raises:
            PDFMergeError: If the merge operation fails
        """
        try:
            logging.info(f"Starting PDF merge with strategy '{strategy}': {input_path} -> {output_path}")
            logging.info(f"Using letterhead: {self.letterhead_path}")
            
            metadata = self.get_doc_info(input_path)
            write_context = self.create_output_context(output_path, metadata)
            read_pdf = self.create_pdf_document(input_path)
            letterhead_pdf = self.create_pdf_document(self.letterhead_path)

            if not all([write_context, read_pdf, letterhead_pdf]):
                error_msg = "Failed to create PDF context or load PDFs"
                logging.error(error_msg)
                raise PDFMergeError(error_msg)

            num_pages = CoreGraphics.CGPDFDocumentGetNumberOfPages(read_pdf)
            logging.info(f"Processing {num_pages} pages")
            
            # Process each page of the content document
            for page_num in range(1, num_pages + 1):
                logging.info(f"Processing page {page_num}")
                page = CoreGraphics.CGPDFDocumentGetPage(read_pdf, page_num)
                letterhead_page = CoreGraphics.CGPDFDocumentGetPage(letterhead_pdf, 1)
                
                if not page or not letterhead_page:
                    error_msg = f"Failed to get page {page_num}"
                    logging.error(error_msg)
                    raise PDFMergeError(error_msg)
                
                media_box = CoreGraphics.CGPDFPageGetBoxRect(page, CoreGraphics.kCGPDFMediaBox)
                if CoreGraphics.CGRectIsEmpty(media_box):
                    media_box = None
                
                CoreGraphics.CGContextBeginPage(write_context, media_box)
                
                # Apply the selected merge strategy
                if strategy == "multiply":
                    self._strategy_multiply(write_context, page, letterhead_page)
                elif strategy == "transparency":
                    self._strategy_transparency(write_context, page, letterhead_page)
                elif strategy == "reverse":
                    self._strategy_reverse(write_context, page, letterhead_page)
                elif strategy == "overlay":
                    self._strategy_overlay(write_context, page, letterhead_page)
                elif strategy == "darken":
                    self._strategy_darken(write_context, page, letterhead_page)
                else:
                    # Default strategy - darken blend mode
                    self._strategy_darken(write_context, page, letterhead_page)
                
                CoreGraphics.CGContextEndPage(write_context)
            
            CoreGraphics.CGPDFContextClose(write_context)
            logging.info("PDF merge completed successfully")

        except Exception as e:
            error_msg = f"Error merging PDFs: {str(e)}"
            logging.error(error_msg, exc_info=True)
            raise PDFMergeError(error_msg)

    def _strategy_multiply(self, context, content_page, letterhead_page):
        """
        Original strategy - draw letterhead first with multiply blend, then content
        
        This is similar to the original implementation but with more explicit control.
        """
        # Draw letterhead first with multiply blend mode
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeMultiply)
        CoreGraphics.CGContextDrawPDFPage(context, letterhead_page)
        CoreGraphics.CGContextRestoreGState(context)
        
        # Draw content on top with normal blend mode
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeNormal)
        CoreGraphics.CGContextDrawPDFPage(context, content_page)
        CoreGraphics.CGContextRestoreGState(context)

    def _strategy_reverse(self, context, content_page, letterhead_page):
        """
        Reverse strategy - draw content first, then letterhead on top with blend mode
        
        This reverses the drawing order which may help with visibility of the letterhead.
        """
        # Draw content first
        CoreGraphics.CGContextDrawPDFPage(context, content_page)
        
        # Draw letterhead on top with multiply blend mode and partial transparency
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeMultiply)
        CoreGraphics.CGContextSetAlpha(context, 0.8)  # 80% opacity
        CoreGraphics.CGContextDrawPDFPage(context, letterhead_page)
        CoreGraphics.CGContextRestoreGState(context)

    def _strategy_transparency(self, context, content_page, letterhead_page):
        """
        Transparency strategy - use transparency layer for better blending
        
        Creates a transparency group which may provide better results for complex blending.
        """
        # Begin transparency layer
        CoreGraphics.CGContextBeginTransparencyLayer(context, None)
        
        # Draw letterhead
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeNormal)
        CoreGraphics.CGContextDrawPDFPage(context, letterhead_page)
        CoreGraphics.CGContextRestoreGState(context)
        
        # Draw content with suitable blend mode
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeOverlay)
        CoreGraphics.CGContextDrawPDFPage(context, content_page)
        CoreGraphics.CGContextRestoreGState(context)
        
        # End transparency layer
        CoreGraphics.CGContextEndTransparencyLayer(context)

    def _strategy_overlay(self, context, content_page, letterhead_page):
        """
        Overlay strategy - uses overlay blend mode for better visibility
        
        The overlay blend mode often works well for watermarks and letterheads as it
        preserves contrast while allowing the letterhead to show through.
        """
        # Draw letterhead first
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetAlpha(context, 0.9)  # Slightly transparent
        CoreGraphics.CGContextDrawPDFPage(context, letterhead_page)
        CoreGraphics.CGContextRestoreGState(context)
        
        # Draw content with overlay blend mode
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeOverlay)
        CoreGraphics.CGContextDrawPDFPage(context, content_page)
        CoreGraphics.CGContextRestoreGState(context)

    def _strategy_darken(self, context, content_page, letterhead_page):
        """
        Darken strategy - uses darken blend mode which can work well for light letterheads
        
        The darken blend mode only keeps pixels that are darker than the background,
        which can work well for dark text/logos on light backgrounds.
        """
        # Draw content first
        CoreGraphics.CGContextDrawPDFPage(context, content_page)
        
        # Draw letterhead with darken blend mode
        CoreGraphics.CGContextSaveGState(context)
        CoreGraphics.CGContextSetBlendMode(context, CoreGraphics.kCGBlendModeDarken)
        CoreGraphics.CGContextDrawPDFPage(context, letterhead_page)
        CoreGraphics.CGContextRestoreGState(context)

    def create_pdf_document(self, path: str) -> Optional[CoreGraphics.CGPDFDocumentRef]:
        """Create PDF document from path"""
        logging.info(f"Creating PDF document from: {path}")
        path_bytes = path.encode('utf-8')
        url = CoreGraphics.CFURLCreateFromFileSystemRepresentation(
            CoreGraphics.kCFAllocatorDefault,
            path_bytes,
            len(path_bytes),
            False
        )
        if not url:
            error_msg = f"Failed to create URL for path: {path}"
            logging.error(error_msg)
            raise PDFMergeError(error_msg)
        doc = CoreGraphics.CGPDFDocumentCreateWithURL(url)
        if not doc:
            error_msg = f"Failed to create PDF document from: {path}"
            logging.error(error_msg)
            raise PDFMergeError(error_msg)
        return doc

    def create_output_context(self, path: str, metadata: Dict[str, Any]) -> Optional[CoreGraphics.CGContextRef]:
        """Create PDF context for output"""
        logging.info(f"Creating output context for: {path}")
        path_bytes = path.encode('utf-8')
        url = CoreGraphics.CFURLCreateFromFileSystemRepresentation(
            CoreGraphics.kCFAllocatorDefault,
            path_bytes,
            len(path_bytes),
            False
        )
        if not url:
            error_msg = f"Failed to create output URL for path: {path}"
            logging.error(error_msg)
            raise PDFMergeError(error_msg)
        context = CoreGraphics.CGPDFContextCreateWithURL(url, None, metadata)
        if not context:
            error_msg = f"Failed to create PDF context for: {path}"
            logging.error(error_msg)
            raise PDFMergeError(error_msg)
        return context

    def get_doc_info(self, file_path: str) -> Dict[str, Any]:
        """Get PDF metadata"""
        logging.info(f"Getting document info from: {file_path}")
        pdf_url = NSURL.fileURLWithPath_(file_path)
        pdf_doc = PDFKit.PDFDocument.alloc().initWithURL_(pdf_url)
        if not pdf_doc:
            error_msg = f"Failed to read PDF metadata from: {file_path}"
            logging.error(error_msg)
            raise PDFMergeError(error_msg)
        
        metadata = pdf_doc.documentAttributes()
        if "Keywords" in metadata:
            keys = metadata["Keywords"]
            mutable_metadata = metadata.mutableCopy()
            mutable_metadata["Keywords"] = tuple(keys)
            return mutable_metadata
        return metadata
