text_processor.py

"""
Text Processing and Repair
OCR text cleanup and repair functions
"""

import re
from datetime import datetime
from openai import OpenAI
from config import config
from logger import ProcessingLogger


class TextRepair:
    """Basic text repair for common OCR errors."""
    BASIC_FIXES = {'rn': 'm', 'vv': 'w', 'cl': 'd',
                   '½': '1/2', '¼': '1/4', '¾': '3/4'}
    FINANCIAL_FIXES = {
        r'\$005,(\d)': r'$\1,500', r'\$000,(\d+)': r'$\1,000'}

class ContentFormatter:
    """Format extracted content into clean markdown."""

    def __init__(self, logger: ProcessingLogger):
        self.logger = logger

        self.client = OpenAI(
            api_key=config.openai_api_key
        )
    
    @staticmethod
    def basic_cleanup(text: str) -> str:
        """Basic text cleanup without AI processing."""
        if not text:
            return ""
        
        # Remove excessive whitespace
        text = re.sub(r'\s+', ' ', text)
        
        # Remove non-printable characters
        text = ''.join(char for char in text if char.isprintable() or char in '\n\r\t')
        
        # Normalize line breaks
        text = re.sub(r'\r\n|\r', '\n', text)
        
        # Remove excessive blank lines
        text = re.sub(r'\n{3,}', '\n\n', text)
        
        return text.strip()

    def build_document_header(self, document_title: str) -> str:
        """Build document header with title and metadata."""
        header = f"# {document_title}\n\n"
        header += f"**Processed:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        return header
    
    def format_content(self, text: str, page_no: int, context=None) -> str:
        """Format content with AI if available, otherwise return cleaned text."""
        try:
            # If no API key, just return cleaned text
            if not config.openai_api_key:
                self.logger.log_step(f"Page {page_no}: Using basic formatting (no API key)")
                return f"## Page {page_no}\n\n{ContentFormatter.basic_cleanup(text)}"
            
            self.logger.log_step(f"Formatting page {page_no} ({len(text)} chars) - this may take 30-90 seconds...")

            prompt = f"""Transform the extracted text using these formatting rules:

* Follow strict Markdown syntax: ## for sections, ### for subsections.
* Only use markdown headers for actual document headings, identified by line breaks, topic changes, word count (less than five words on average), and font size or weight.
* You are permitted to remove page numbers and other document metadata.
* Ignore page breaks. Retain header continuity across pages.
* Ignore page breaks. Preserve table formatting across pages.

### Checkbox Detection and Formatting

* Scan ALL text for checkbox indicators, not just row beginnings
* Recognize these checkbox patterns as CHECKED:
  - [x] or [X] (bracketed X)
  - ☑ or ✓ or ✔ (checkbox symbols)
  - (x) or (X) (parenthetical X)
  - Standalone X or x when clearly part of a checkbox field
  
* Recognize these patterns as UNCHECKED:
  - [ ] (empty brackets)
  - ☐ (empty checkbox symbol)
  - ( ) (empty parentheses)
  - O or o when clearly part of a checkbox field

* Format checkbox fields as:
  - **Checked**:  *[SELECTED]* **[Field Label]**: [Additional field value if any]
  - **Unchecked**: **[Field Label]**: [Additional field value if any]

## Transforming Tables

When you encounter data that appears to be in table format (aligned columns of data):

* Treat text that wraps within cells as a single cell.
* Identify column headers (typically the utmost top labels)
* Identify row headers (typically the leftmost column)
* For each cell value, create a **flattened structured entry** that connects it to both its row and column, using this format:
    * **[Column Label, including wrapped cell content]** - **[Row Label, including wrapped cell content]**: [Value]

## Formatting principles

* Preserve bullet points and numbered lists
* Maintain logical groupings of related information
* Keep indentation patterns and data hierarchy

## Preservation rules

* Preserve all original text exactly as-is
* Add annotations in a separate notation, e.g., *[SELECTED]*


{text}
"""
            
            response = self.client.chat.completions.create(
                model=config.openai_model,
                messages=[
                    {
                        "role": "system", 
                        "content": "You are an expert document analysis AI. Format text as clean Markdown. Preserve all content exactly."
                    },
                    {
                        "role": "user", 
                        "content": prompt
                    }
                ],
                temperature=getattr(config, 'temperature', 0.1),
                max_completion_tokens=getattr(config, 'max_output_tokens', 16384),  # Use config value
                timeout=120  # 2 minute timeout for complex formatting
            )

            result = response.choices[0].message.content.strip()
            self.logger.log_success(f"Page {page_no} formatting completed - {len(result)} chars output")
            return result
        
        except Exception as e:
            self.logger.log_error(f"Page {page_no} formatting failed: {type(e).__name__} - {e}")
            # Return cleaned text with page header as fallback
            return f"## Page {page_no}\n\n{ContentFormatter.basic_cleanup(text)}"