# AWS SDK for Python — provides S3 client for object storage operations
import boto3
# Standard library JSON module — serialises job dictionaries into message payloads
import json
# Enum base class — used for defining document type classification constants later
from enum import Enum

# Create a reusable S3 client; single instance reused across all invocations
s3 = boto3.client("s3")

# Entry point triggered by an S3 upload event (Lambda or webhook)
def handle_upload(event):
    # S3 event notifications can batch multiple object creations; iterate each
    for record in event["Records"]:
        # Extract the source bucket name from the S3 event notification metadata
        bucket = record["s3"]["bucket"]["name"]
        # Extract the full object key (path) of the uploaded PDF file
        key = record["s3"]["object"]["key"]

        # Build a lightweight job descriptor for downstream workers
        job = {
            # Derive unique job ID from the PDF filename, stripping the .pdf extension
            "job_id": key.split("/")[-1].replace(".pdf", ""),
            # Preserve the bucket so workers know where to download the PDF from
            "bucket": bucket,
            # Preserve the full S3 key so workers can fetch the exact object
            "key": key,
            # Record file size for monitoring dashboards and cost tracking
            "size": record["s3"]["object"]["size"],
            # Extract tenant/organisation from the first path segment of the key
            "tenant": key.split("/")[0],
        }

        # Push the job onto the message broker for asynchronous consumption
        # Publish to RabbitMQ
        channel.basic_publish(
            # Default exchange routes messages directly by the queue name provided
            exchange="",
            # Queue name that ingestion workers are bound to consume from
            routing_key="pdf-ingest",
            # Serialise the job dictionary as a JSON string for AMQP transport
            body=json.dumps(job),
        )

# Import PyPDF2 to read and examine PDF file contents programmatically
import PyPDF2

# Classify a PDF as "digital" (has extractable text) or "scanned" (image-only pages)
def classify_pdf(pdf_path: str) -> str:
    # Open the PDF file in binary read mode
    with open(pdf_path, "rb") as f:
        # Create a PdfReader instance to parse the PDF structure
        reader = PyPDF2.PdfReader(f)
        # Iterate through every page in the document
        for page in reader.pages:
            # Attempt to extract any text content from the current page
            text = page.extract_text()
            # If non-whitespace text is found, the PDF is digital
            if text.strip():
                return "digital"  # has extractable text
    # No extractable text found on any page — must be a scanned image PDF
    return "scanned"  # all pages are images

# Regular expression module — used to extract structured fields from raw text
import re
# Typing import for type hinting the return value as a dictionary
from typing import Dict

# Extract structured fields based on document type using regex patterns
def extract_fields(text: str, doc_type: str) -> Dict:
    # Invoice-specific extraction patterns
    if doc_type == "invoice":
        return {
            # Capture the invoice number after "Invoice" or "Invoice:" (non-whitespace token)
            "invoice_number": re.search(r"Invoice[:\s]*(\S+)", text).group(1) if re.search(r"Invoice[:\s]*(\S+)", text) else None,
            # Capture the date string after "Date" or "Date:"
            "date": re.search(r"Date[:\s]*(\S+)", text).group(1) if re.search(r"Date[:\s]*(\S+)", text) else None,
            # Capture the total amount (optional $ prefix, digits, commas, decimal)
            "total": re.search(r"Total[:\s]*\$?([\d,]+\.?\d*)", text).group(1) if re.search(r"Total[:\s]*\$?([\d,]+\.?\d*)", text) else None,
        }
    # Contract-specific extraction patterns
    elif doc_type == "contract":
        return {
            # Find all party names following "Between" or "Party:" (capitalised words)
            "parties": re.findall(r"(?:Between|Party)[:\s]*([A-Z][A-Za-z\s]+)", text),
            # Capture the effective date after "Effective Date" or "Effective Date:"
            "effective_date": re.search(r"Effective[:\s]*Date[:\s]*(\S+)", text).group(1) if re.search(r"Effective[:\s]*Date[:\s]*(\S+)", text) else None,
        }
    # Extend with additional document type handlers as needed
    # Add more document types
    # Return empty dict for unrecognised document types
    return {}

# AWS SDK — provides the Textract client for managed OCR and document analysis
import boto3

# Create a Textract client; AWS's managed service handles OCR at scale
textract = boto3.client("textract")

# Use AWS Textract to extract text from a digital PDF stored in S3
def fast_extract(s3_bucket: str, s3_key: str) -> Dict:
    # Call detect_document_text — synchronous API for text detection
    response = textract.detect_document_text(
        # Reference the PDF directly in S3 instead of streaming the bytes
        Document={"S3Object": {"Bucket": s3_bucket, "Name": s3_key}}
    )
    # Concatenate all detected text lines (BlockType == "LINE") separated by newlines
    text = "\n".join(
        block["Text"] for block in response["Blocks"]
        if block["BlockType"] == "LINE"
    )
    # Return the full extracted text and total block count (rough page estimate)
    return {"text": text, "pages": len(response["Blocks"])}

# Thread pooling utilities — run OCR on multiple pages concurrently
from concurrent.futures import ThreadPoolExecutor, as_completed
# Python bindings for Tesseract OCR engine — converts image pixels to text
import pytesseract
# Converts PDF pages into PIL Image objects at configurable DPI
from pdf2image import convert_from_path

# OCR a single page image with Tesseract (I/O-bound, fits threading well)
def ocr_page(page_img) -> str:
    # image_to_string runs Tesseract on the PIL image, English language model
    return pytesseract.image_to_string(page_img, lang="eng")

# OCR all pages of a scanned PDF in parallel using a thread pool
def ocr_pdf_parallel(pdf_path: str, max_workers: int = 8) -> str:
    # Convert every PDF page to a high-resolution PIL image (300 DPI)
    images = convert_from_path(pdf_path, dpi=300)
    # Pre-allocate a list to hold results in original page order
    pages = [None] * len(images)

    # Thread pool executes OCR calls concurrently; 8 workers matches typical CPU core count
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit each page for OCR; map each Future back to its page index
        futures = {executor.submit(ocr_page, img): i for i, img in enumerate(images)}
        # Process results as each Future completes (not necessarily in order)
        for future in as_completed(futures):
            # Retrieve the original page index from the futures mapping
            idx = futures[future]
            # Store the OCR result at the correct index to preserve page order
            pages[idx] = future.result()

    # Join all pages with double newlines to reconstruct the full document text
    return "\n\n".join(pages)

# DonutProcessor handles image loading, tokenisation, and decoding for the Donut model
from transformers import DonutProcessor, VisionEncoderDecoderModel
# PyTorch — tensor operations and GPU management for model inference
import torch

# Load the pre-trained Donut processor (image transforms + tokeniser) from HuggingFace Hub
processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base")
# Load the pre-trained vision-encoder-decoder model (Donut) from HuggingFace Hub
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")

# Extract structured JSON directly from a scanned document image (no separate OCR step)
def extract_from_image(image_path: str) -> Dict:
    # Load the image file as a PIL Image using the Donut processor's loader
    image = processor.load_image(image_path)
    # Preprocess the image into pixel tensors; "pt" means PyTorch tensors
    pixel_values = processor(image, return_tensors="pt").pixel_values

    # Run the encoder-decoder model to generate a text token sequence
    outputs = model.generate(
        pixel_values,
        max_length=512,    # Cap output token length to bound inference time
        early_stopping=True,  # Stop generation when end-of-sequence token is emitted
    )

    # Decode the first (and only) beam's token IDs back to a text string
    sequence = processor.batch_decode(outputs)[0]
    # Convert the encoded JSON string into a Python dictionary (structured output)
    return processor.token2json(sequence)

# OpenAI Python client — interfaces with GPT models for structured extraction
import openai
# Tenacity provides decorator-based retry logic with exponential backoff
from tenacity import retry, stop_after_attempt, wait_exponential

# Create the OpenAI client singleton; reads API key from environment by default
client = openai.OpenAI()

# Retry up to 3 times with exponential backoff (2s, 4s, 10s) for transient API failures
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
def llm_extract(text: str, doc_type: str) -> Dict:
    # Call the chat completions endpoint with JSON mode enabled
    response = client.chat.completions.create(
        model="gpt-4o-mini",               # Cost-effective model with good extraction accuracy
        response_format={"type": "json_object"},  # Force valid JSON in the response
        messages=[
            {
                "role": "system",
                # System prompt instructs the LLM on the expected JSON schema
                "content": f"Extract structured data from this {doc_type}. Return JSON with fields: invoice_number, date, total, vendor, customer. If a field is missing, set it to null.",
            },
            # Truncate input text to 8000 characters to stay within the model's context window
            {"role": "user", "content": text[:8000]},
        ],
    )
    # Parse the model's JSON string response into a Python dictionary
    return json.loads(response.choices[0].message.content)

# Pre-warmed model pool — loads multiple model copies into GPU memory at startup
class ModelPool:
    # pool_size controls concurrency: 4 copies handle 4 simultaneous inference requests
    def __init__(self, model_name: str, device: str = "cuda", pool_size: int = 4):
        self.models = []    # Will hold separate model instances for concurrent access
        for _ in range(pool_size):
            model = load_model(model_name)        # Load model weights from disk or cache
            model.to(device)                       # Move model to GPU (cuda) or CPU
            model.eval()                           # Set to evaluation mode (disables dropout etc.)
            self.models.append(model)              # Append to the pool

    # Dispatch an inference request to the next available model copy
    def infer(self, inputs):
        # Round-robin across models for concurrent requests
        model = self.models[self.counter % len(self.models)]  # Select model by round-robin index
        self.counter += 1                                      # Advance the counter for next call
        with torch.no_grad():                                  # Disable gradient tracking for speed
            return model(inputs)                               # Run forward pass

# Process multiple page images in a single GPU batch to maximise throughput
def batch_ocr(images, batch_size=8):
    results = []   # Accumulator for decoded text from each image
    # Iterate over the image list in fixed-size chunks
    for i in range(0, len(images), batch_size):
        batch = images[i:i + batch_size]                        # Slice the current batch of images
        inputs = processor(batch, return_tensors="pt", padding=True).to("cuda")  # Preprocess and pad batch to GPU

        with torch.no_grad():                                   # Disable gradient computation for inference
            outputs = model.generate(**inputs, max_length=512)   # Generate text tokens for the entire batch

        # Decode each generated token sequence into a readable string
        for out in outputs:
            results.append(processor.decode(out, skip_special_tokens=True))  # Strip padding/sep tokens

    return results  # Return all decoded page texts in original order

# FastAPI — async web framework for high-concurrency API endpoints
from fastapi import FastAPI, BackgroundTasks, UploadFile
# uuid4 generates unique job IDs without a centralised counter (no collision risk)
from uuid import uuid4

# Initialise the FastAPI application
app = FastAPI()
# In-memory dictionary as placeholder for job results; Redis or similar in production
results_cache = {}  # In production, use Redis

# POST endpoint that accepts a PDF upload and returns immediately with a job ID
@app.post("/extract")
async def extract_pdf(file: UploadFile):
    job_id = str(uuid4())        # Generate a unique identifier for this extraction job
    content = await file.read()  # Read the entire uploaded file as bytes (non-blocking)

    # Start processing in background
    background_tasks.add_task(process_pdf, job_id, content)  # Enqueue async task; does not block response

    # Return immediately with a job ID
    return {"job_id": job_id, "status": "processing"}  # Client polls this ID for completion

# GET endpoint that the client polls to check job status and retrieve results
@app.get("/results/{job_id}")
async def get_results(job_id: str):
    result = results_cache.get(job_id)  # Look up cached result; None if not yet complete
    if not result:
        return {"status": "pending"}    # Job still in-flight — tell the client to try again
    return {"status": "done", "data": result}  # Job finished — return extracted data

# Main async worker function that orchestrates the entire document processing pipeline
async def process_pdf(job_id: str, content: bytes, doc_type: str):
    start = time.time()  # Record start time for per-document latency tracking

    # Step 1: Write to temp file
    tmp = f"/tmp/{job_id}.pdf"  # Temporary file path unique to this job
    with open(tmp, "wb") as f:  # Open temp file in binary write mode
        f.write(content)        # Persist the raw PDF bytes to disk for library access

    # Step 2: Classify
    pdf_type = classify_pdf(tmp)  # Determine if PDF is digital text or scanned image

    # Step 3: Extract
    # Route to the cheapest adequate extraction method based on PDF type
    if pdf_type == "digital":
        text = extract_text_from_pdf(tmp)   # Direct text extraction — no OCR needed
        result = extract_fields(text, doc_type)  # Regex-based field extraction from plain text
    elif pdf_type == "scanned":
        text = ocr_pdf(tmp)                 # OCR fallback — extracts raw text from images
        result = extract_with_donut(tmp)    # GPU batch — structured JSON direct from image
    else:
        result = llm_extract(text, doc_type)  # LLM fallback for complex/unusual documents

    # Step 4: Validate
    result["confidence"] = validate_extraction(result)  # Score extraction quality (0.0–1.0)

    # Step 5: Store
    s3.put_object(                          # Persist the structured result to S3
        Bucket="processed-docs",             # Destination bucket for completed extractions
        Key=f"{doc_type}/{job_id}.json",     # Organise results by doc type and job ID
        Body=json.dumps(result),             # Serialise the result dict as a JSON string
    )

    # Step 6: Cache for fast retrieval
    await redis.set(f"result:{job_id}", json.dumps(result), ex=3600)  # Cache with 1-hour TTL

    elapsed = time.time() - start
    logger.info(f"Processed {job_id} in {elapsed:.2f}s (type={pdf_type})")  # Log latency per document

    os.remove(tmp)  # Clean up the temporary file to avoid disk filling

# Custom exception for transient failures that can be safely retried
class RetryableError(Exception):
    pass

# Custom exception for permanent failures that should not be retried
class NonRetryableError(Exception):
    pass

# Centralised failure handler — decides whether to retry or dead-letter
def handle_failure(job_id, error, pdf_path):
    # Check for non-retryable errors first (password-protected, corrupt, etc.)
    if isinstance(error, NonRetryableError):
        # Password protected or corrupted — dead letter queue
        dlq.send(job_id, {"reason": str(error)})  # Send to dead-letter queue for manual review
        return                                      # No further processing for this job

    # Check for retryable errors (network timeouts, S3 throttling, etc.)
    if isinstance(error, RetryableError):
        # Transient error — retry with backoff
        retry_count = get_retry_count(job_id)       # Fetch number of previous attempts from Redis
        if retry_count < 3:                          # Maximum 3 retries before giving up
            queue.send_with_delay(job_id, delay=2 ** retry_count * 10)  # Exponential backoff: 10s, 20s, 40s
        else:
            dlq.send(job_id, {"reason": "max retries exceeded"})  # Dead-letter after exhausting retries