import requests
from io import BytesIO
import PyPDF2
from datetime import datetime

@data_loader
def load_pdf_from_github(**kwargs):

    # Load PDF from GitHub repository
    github_pdf_url = "<https://raw.githubusercontent.com/mage-ai/datasets/master/great_attractor_pdf.pdf>"
    
    print(f"Fetching PDF from GitHub: {github_pdf_url}")
    
    # Download the PDF
    try:
        response = requests.get(github_pdf_url)
        response.raise_for_status()  # Ensure we got a valid response
        
        # Read the PDF content
        pdf_content = BytesIO(response.content)
        pdf_reader = PyPDF2.PdfReader(pdf_content)
        
        # Extract metadata
        pdf_info = {
            "title": "Evidence of the Great Attractor and Great Repeller",
            "author": "Christopher C. O'Neill",
            "num_pages": len(pdf_reader.pages),
            "source_url": github_pdf_url,
            "fetch_time": datetime.now().isoformat()
        }
        
        # Extract text from all pages
        all_text = ""
        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            page_text = page.extract_text()
            all_text += page_text + "\\n\\n"
        
        print(f"Successfully extracted {pdf_info['num_pages']} pages from PDF")
        
        # Create result with both metadata and content
        result = {
            "metadata": pdf_info,
            "content": all_text
        }
        
        return result
    except Exception as e:
        print(f"Error fetching PDF from GitHub: {e}")
        # Instead of providing sample text, let's throw an error to ensure we don't proceed without the actual document
        raise Exception(f"Failed to load PDF document: {e}")

import re

@transformer
def clean_document(data, **kwargs):
    """
    Clean and format the extracted PDF text
    """
    # Check if data is available
    if data is None or "content" not in data:
        raise Exception("No PDF content available to clean")
    
    # Extract text and metadata from previous block
    text = data["content"]
    metadata = data["metadata"]
    
    print("Cleaning document text...")
    
    # Clean up common PDF extraction issues
    # 1. Replace multiple newlines with a single newline
    text = re.sub(r'\\n{3,}', '\\n\\n', text)
    
    # 2. Fix broken words that might have been split across lines
    text = re.sub(r'(\\w+)-\\n(\\w+)', r'\\1\\2', text)
    
    # 3. Remove headers and footers that repeat on each page
    text = re.sub(r'Journal of High Energy Physics,\\s+Gravitation and Cosmology.+?\\d+', '', text)
    text = re.sub(r'DOI:.+?\\d+', '', text)
    
    # 4. Normalize whitespace
    text = re.sub(r' {2,}', ' ', text)
    text = text.strip()
    
    # Add document information at the beginning
    title_header = f"# {metadata['title']}\\n"
    author_header = f"## By {metadata['author']}\\n"
    source_header = f"Source: {metadata['source_url']}\\n\\n"
    
    formatted_text = title_header + author_header + source_header + text
    
    # Return cleaned document with metadata
    result = {
        "text": formatted_text,
        "metadata": metadata,
        "character_count": len(formatted_text),
        "processing_time": kwargs.get("execution_date", "Unknown")
    }
    
    print(f"Document cleaned: {len(formatted_text)} characters")
    return result

@transformer
def chunk_document(data, **kwargs):
    """
    Split the cleaned document into manageable chunks with overlap
    """
    # Check if data is available
    if data is None or "text" not in data:
        raise Exception("No cleaned document text available to chunk")

    text = data["text"]
    metadata = data["metadata"]
    
    print("Chunking document with fixed size approach...")
    
    # Set chunking parameters
    chunk_size = 4000  # characters per chunk
    overlap = 500      # overlap between chunks to maintain context
    
    # Calculate how many chunks we'll need
    doc_length = len(text)
    estimated_chunks = (doc_length // (chunk_size - overlap)) + 1
    print(f"Document length: {doc_length} characters")
    print(f"Will create approximately {estimated_chunks} chunks")
    
    # Create chunks with simple fixed-size approach
    chunks = []
    
    # Special case: if text is smaller than chunk_size, just use one chunk
    if doc_length <= chunk_size:
        chunks.append({
            "chunk_id": 0,
            "text": text,
            "start_char": 0,
            "end_char": doc_length,
            "character_count": doc_length,
            "doc_id": metadata['source_url']
        })
    else:
        # Use a straightforward loop with explicit indices
        chunk_id = 0
        for i in range(0, doc_length, chunk_size - overlap):
            # Calculate chunk boundaries
            start = i
            end = min(i + chunk_size, doc_length)
            
            # Don't create tiny chunks at the end
            if end - start < 200 and chunk_id > 0:
                break
                
            # Extract chunk text
            chunk_text = text[start:end]
            
            # Create chunk with metadata
            chunks.append({
                "chunk_id": chunk_id,
                "text": chunk_text,
                "start_char": start,
                "end_char": end,
                "character_count": len(chunk_text),
                "doc_id": metadata['source_url']
            })
            
            chunk_id += 1
            
            # Log progress for larger documents
            if chunk_id % 5 == 0:
                print(f"Created {chunk_id} chunks so far...")
    
    # Create result
    result = {
        "chunks": chunks,
        "chunk_count": len(chunks),
        "total_characters": doc_length,
        "metadata": metadata
    }
    
    print(f"Document processed into {len(chunks)} chunks")
    return result