import os
import PyPDF2
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
import io
import boto3
import tempfile
import shutil
from botocore.exceptions import ClientError, NoCredentialsError
from typing import Dict, List, Optional, Tuple

class S3PDFMerger:
    """PDF Merger that works with S3-compatible storage."""
    
    def __init__(self, s3_config: Dict):
        """
        Initialize S3 PDF Merger with configuration.
        
        Args:
            s3_config (Dict): S3 configuration containing:
                - access_key_id: AWS access key
                - secret_access_key: AWS secret key  
                - region: AWS region
                - endpoint: S3 endpoint URL
                - default_bucket: Default bucket name
        """
        self.s3_config = s3_config
        self.s3_client = self._create_s3_client()
        self.temp_dir = None
        
    def _create_s3_client(self):
        """Create and configure S3 client."""
        try:
            return boto3.client(
                's3',
                aws_access_key_id=self.s3_config['access_key_id'],
                aws_secret_access_key=self.s3_config['secret_access_key'],
                region_name=self.s3_config['region'],
                endpoint_url=self.s3_config['endpoint']
            )
        except Exception as e:
            raise ConnectionError(f"Failed to create S3 client: {e}")
    
    def list_pdf_files(self, bucket_name: str, prefix: str = "") -> List[str]:
        """
        List all PDF files in S3 bucket with optional prefix.
        
        Args:
            bucket_name (str): S3 bucket name
            prefix (str): Optional prefix to filter files
            
        Returns:
            List[str]: List of PDF file keys
        """
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=bucket_name,
                Prefix=prefix
            )
            
            if 'Contents' not in response:
                return []
            
            pdf_files = [
                obj['Key'] for obj in response['Contents']
                if obj['Key'].lower().endswith('.pdf') and obj['Size'] > 0
            ]
            
            return sorted(pdf_files)
            
        except ClientError as e:
            print(f"Error listing files from S3: {e}")
            return []
    
    def download_file(self, bucket_name: str, key: str, local_path: str) -> bool:
        """
        Download a file from S3 to local temporary storage.
        
        Args:
            bucket_name (str): S3 bucket name
            key (str): S3 object key
            local_path (str): Local file path to save
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.s3_client.download_file(bucket_name, key, local_path)
            return True
        except ClientError as e:
            print(f"Error downloading {key}: {e}")
            return False
    
    def upload_file(self, local_path: str, bucket_name: str, key: str) -> bool:
        """
        Upload a file from local storage to S3.
        
        Args:
            local_path (str): Local file path
            bucket_name (str): S3 bucket name
            key (str): S3 object key
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.s3_client.upload_file(local_path, bucket_name, key)
            return True
        except ClientError as e:
            print(f"Error uploading {key}: {e}")
            return False
    
    def get_file_size(self, bucket_name: str, key: str) -> int:
        """
        Get the size of a file in S3.
        
        Args:
            bucket_name (str): S3 bucket name
            key (str): S3 object key
            
        Returns:
            int: File size in bytes, 0 if error
        """
        try:
            response = self.s3_client.head_object(Bucket=bucket_name, Key=key)
            return response['ContentLength']
        except ClientError:
            return 0

    def merge_pdfs_from_s3(self, bucket_name: str, prefix: str = "", 
                          output_key: str = "merged_document.pdf", 
                          add_page_numbers: bool = True) -> Tuple[bool, Dict]:
        """
        Merge all PDF files from S3 bucket with an index table.
        
        Args:
            bucket_name (str): S3 bucket name
            prefix (str): Optional prefix to filter files (like folder path)
            output_key (str): S3 key for the output merged PDF
            add_page_numbers (bool): Whether to add visible page numbers
            
        Returns:
            Tuple[bool, Dict]: (success, results_info)
        """
        
        # Create temporary directory for processing
        self.temp_dir = tempfile.mkdtemp()
        
        try:
            print(f"🔍 Searching for PDF files in s3://{bucket_name}/{prefix}")
            
            # List PDF files in S3
            pdf_files = self.list_pdf_files(bucket_name, prefix)
            
            if not pdf_files:
                return False, {"error": "No PDF files found in S3 bucket"}
            
            print(f"📁 Found {len(pdf_files)} PDF files:")
            for pdf_file in pdf_files:
                size = self.get_file_size(bucket_name, pdf_file)
                print(f"   • {pdf_file} ({self._format_file_size(size)})")
            
            # Download all PDF files to temporary directory
            print(f"\n⬇️  Downloading PDF files...")
            local_files = []
            
            for i, pdf_key in enumerate(pdf_files, 1):
                local_filename = f"pdf_{i:03d}_{os.path.basename(pdf_key)}"
                local_path = os.path.join(self.temp_dir, local_filename)
                
                print(f"   📥 Downloading {pdf_key}...")
                if self.download_file(bucket_name, pdf_key, local_path):
                    local_files.append({
                        'local_path': local_path,
                        'original_key': pdf_key,
                        'display_name': self._get_display_name(pdf_key)
                    })
                else:
                    print(f"   ❌ Failed to download {pdf_key}")
            
            if not local_files:
                return False, {"error": "Failed to download any PDF files"}
            
            # Process files for merging
            document_info = []
            current_page = 2  # Start from page 2 (index will be page 1)
            
            print(f"\n📄 Processing PDF files for merger...")
            
            for file_info in local_files:
                try:
                    with open(file_info['local_path'], 'rb') as file:
                        pdf_reader = PyPDF2.PdfReader(file)
                        num_pages = len(pdf_reader.pages)
                        
                        document_info.append({
                            'name': file_info['display_name'],
                            'start_page': current_page,
                            'end_page': current_page + num_pages - 1,
                            'file_path': file_info['local_path'],
                            'original_key': file_info['original_key']
                        })
                        
                        current_page += num_pages
                        print(f"   ✅ {file_info['display_name']}: {num_pages} pages")
                        
                except Exception as e:
                    print(f"   ❌ Error processing {file_info['display_name']}: {e}")
                    continue
            
            if not document_info:
                return False, {"error": "No valid PDF files could be processed"}
            
            # Create merged PDF
            print(f"\n🔗 Creating merged PDF...")
            merged_path = os.path.join(self.temp_dir, "merged_output.pdf")
            
            # Create index page
            index_pdf = self._create_index_page(document_info)
            
            # Create final merged PDF
            pdf_writer = PyPDF2.PdfWriter()
            
            # Add index page first
            index_reader = PyPDF2.PdfReader(index_pdf)
            pdf_writer.add_page(index_reader.pages[0])
            
            # Add all other PDFs with optional page numbering
            page_number = 2  # Start numbering from page 2
            
            for doc_info in document_info:
                try:
                    with open(doc_info['file_path'], 'rb') as file:
                        pdf_reader = PyPDF2.PdfReader(file)
                        
                        for page in pdf_reader.pages:
                            if add_page_numbers:
                                numbered_page = self._add_page_number_to_page(page, page_number)
                                pdf_writer.add_page(numbered_page)
                            else:
                                pdf_writer.add_page(page)
                            page_number += 1
                except Exception as e:
                    print(f"   ❌ Error merging {doc_info['name']}: {e}")
                    continue
            
            # Save merged PDF locally
            with open(merged_path, 'wb') as output_file:
                pdf_writer.write(output_file)
            
            # Upload merged PDF to S3
            print(f"\n⬆️  Uploading merged PDF to S3...")
            upload_success = self.upload_file(merged_path, bucket_name, output_key)
            
            if upload_success:
                # Get final file size
                final_size = self.get_file_size(bucket_name, output_key)
                
                results = {
                    "success": True,
                    "output_location": f"s3://{bucket_name}/{output_key}",
                    "total_pages": len(pdf_writer.pages),
                    "documents_merged": len(document_info),
                    "final_size": final_size,
                    "document_summary": document_info
                }
                
                print(f"✅ Merged PDF uploaded successfully!")
                print(f"📍 Location: s3://{bucket_name}/{output_key}")
                print(f"📄 Total pages: {results['total_pages']}")
                print(f"📁 Documents merged: {results['documents_merged']}")
                print(f"💾 Final size: {self._format_file_size(final_size)}")
                
                # Print document summary
                print(f"\n📋 Document Summary:")
                for doc in document_info:
                    if doc['start_page'] == doc['end_page']:
                        pages_text = f"Page {doc['start_page']}"
                    else:
                        pages_text = f"Pages {doc['start_page']}-{doc['end_page']}"
                    print(f"   • {doc['name']}: {pages_text}")
                
                return True, results
            else:
                return False, {"error": "Failed to upload merged PDF to S3"}
                
        except Exception as e:
            return False, {"error": f"Merge operation failed: {str(e)}"}
        finally:
            # Clean up temporary directory
            if self.temp_dir and os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
                self.temp_dir = None

    def _get_display_name(self, s3_key: str) -> str:
        """Extract display name from S3 key."""
        # Get filename without extension
        filename = os.path.basename(s3_key)
        return os.path.splitext(filename)[0]
    
    def _format_file_size(self, size_bytes: int) -> str:
        """Convert bytes to human readable format."""
        if size_bytes < 1024:
            return f"{size_bytes} B"
        elif size_bytes < 1024**2:
            return f"{size_bytes/1024:.1f} KB"
        elif size_bytes < 1024**3:
            return f"{size_bytes/(1024**2):.1f} MB"
        else:
            return f"{size_bytes/(1024**3):.1f} GB"

    def _create_index_page(self, document_info):
        """Create a table of contents page."""
        buffer = io.BytesIO()
        
        # Create canvas
        c = canvas.Canvas(buffer, pagesize=letter)
        width, height = letter
        
        # Title
        c.setFont("Helvetica-Bold", 18)
        title_width = c.stringWidth("Table of Contents", "Helvetica-Bold", 18)
        c.drawString((width - title_width) / 2, height - 60, "Table of Contents")
        
        # Add a line under title
        c.setStrokeColorRGB(0.3, 0.3, 0.3)
        c.setLineWidth(1)
        c.line(50, height - 80, width - 50, height - 80)
        
        # Headers
        c.setFont("Helvetica-Bold", 12)
        c.setFillColorRGB(0, 0, 0)
        c.drawString(60, height - 110, "Document Name")
        c.drawString(width - 150, height - 110, "Pages")
        
        # Draw line under headers
        c.setLineWidth(0.5)
        c.line(50, height - 120, width - 50, height - 120)
        
        # Document entries
        c.setFont("Helvetica", 11)
        y_position = height - 145
        
        for i, doc in enumerate(document_info):
            if y_position < 80:  # If we're running out of space
                c.drawString(60, y_position, "... (additional documents)")
                break
            
            # Alternate row background
            if i % 2 == 0:
                c.setFillColorRGB(0.95, 0.95, 0.95)
                c.rect(50, y_position - 3, width - 100, 18, fill=1, stroke=0)
            
            c.setFillColorRGB(0, 0, 0)
            
            # Document name (truncate if too long)
            doc_name = doc['name']
            if len(doc_name) > 40:
                doc_name = doc_name[:37] + "..."
            c.drawString(60, y_position, doc_name)
            
            # Page range
            if doc['start_page'] == doc['end_page']:
                page_text = str(doc['start_page'])
            else:
                page_text = f"{doc['start_page']}-{doc['end_page']}"
            
            # Right-align page numbers
            page_width = c.stringWidth(page_text, "Helvetica", 11)
            c.drawString(width - 100 - page_width, y_position, page_text)
            
            y_position -= 22
        
        # Add footer with generation info
        c.setFont("Helvetica", 8)
        c.setFillColorRGB(0.5, 0.5, 0.5)
        c.drawString(60, 40, f"Generated from S3 bucket - Total documents: {len(document_info)}")
        
        # Add page number to index page
        c.drawString(width - 80, 40, "Page 1")
        
        c.save()
        buffer.seek(0)
        return buffer

    def _add_page_number_to_page(self, page, page_number):
        """Add a page number to a PDF page."""
        # Create a new PDF with just the page number
        packet = io.BytesIO()
        can = canvas.Canvas(packet, pagesize=letter)
        
        # Get page dimensions
        page_width = float(page.mediabox.width)
        page_height = float(page.mediabox.height)
        
        # Add page number at bottom center
        can.setFont("Helvetica", 10)
        page_text = f"Page {page_number}"
        text_width = can.stringWidth(page_text, "Helvetica", 10)
        
        # Position at bottom center
        x_position = (page_width - text_width) / 2
        y_position = 30
        
        can.setFillColorRGB(0.3, 0.3, 0.3)
        can.drawString(x_position, y_position, page_text)
        can.save()
        
        # Move to the beginning of the StringIO buffer
        packet.seek(0)
        
        # Create a new PDF reader for the page number
        page_number_pdf = PyPDF2.PdfReader(packet)
        page_number_page = page_number_pdf.pages[0]
        
        # Merge the page number with the original page
        page.merge_page(page_number_page)
        
        return page

def merge_s3_pdfs(s3_config: Dict, bucket_name: str, prefix: str = "", 
                  output_key: str = "merged_document.pdf", 
                  add_page_numbers: bool = True) -> Tuple[bool, Dict]:
    """
    Convenience function to merge PDFs from S3 bucket.
    
    Args:
        s3_config (Dict): S3 configuration
        bucket_name (str): S3 bucket name
        prefix (str): Optional prefix to filter files
        output_key (str): S3 key for output file
        add_page_numbers (bool): Whether to add page numbers
        
    Returns:
        Tuple[bool, Dict]: (success, results)
    """
    merger = S3PDFMerger(s3_config)
    return merger.merge_pdfs_from_s3(bucket_name, prefix, output_key, add_page_numbers)

def main():
    """Interactive main function for S3 PDF merging."""
    print("🔗 S3 PDF Merger with Index Generator")
    print("=" * 45)
    
    # Default S3 configuration (can be overridden)
    default_s3_config = {
        "access_key_id": "HWYL2OAMI52ZMGSXC3J2",
        "secret_access_key": "W8MxaVHsundr0cATpN4suZ7RKCK8HKqsjnxqRzqD",
        "region": "ap-south-1",
        "endpoint": "https://ap-south-1.linodeobjects.com",
        "default_bucket": "minaions"
    }
    
    # Get bucket name
    bucket_name = input(f"🪣 Enter S3 bucket name (default: {default_s3_config['default_bucket']}): ").strip()
    if not bucket_name:
        bucket_name = default_s3_config['default_bucket']
    
    # Get prefix (folder path)
    prefix = input("📁 Enter prefix/folder path (press Enter for root): ").strip()
    
    # Get output filename
    output_key = input("📄 Enter output filename (default: merged_document.pdf): ").strip()
    if not output_key:
        output_key = "merged_document.pdf"
    elif not output_key.lower().endswith('.pdf'):
        output_key += '.pdf'
    
    # Add prefix to output if specified
    if prefix and not output_key.startswith(prefix):
        output_key = f"{prefix.rstrip('/')}/{output_key}"
    
    # Ask about page numbering
    add_numbers = input("🔢 Add visible page numbers to each page? (y/n, default: y): ").strip().lower()
    add_page_numbers = add_numbers != 'n'
    
    print(f"\n🚀 Starting S3 PDF merge process...")
    print(f"🪣 Bucket: {bucket_name}")
    print(f"📁 Prefix: {prefix if prefix else '(root)'}")
    print(f"📄 Output: {output_key}")
    
    try:
        # Run the merger
        success, results = merge_s3_pdfs(
            s3_config=default_s3_config,
            bucket_name=bucket_name,
            prefix=prefix,
            output_key=output_key,
            add_page_numbers=add_page_numbers
        )
        
        if not success:
            print(f"\n❌ Merge failed!")
            if 'error' in results:
                print(f"Error: {results['error']}")
                
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")

if __name__ == "__main__":
    main()