#!/usr/bin/env python3
"""
reflow_ocr.py

Reflows OCR'd book text whose lines are broken arbitrarily by the
physical page width, turning it back into normal flowing paragraphs.

Handles:
  - Hyphenated words split across a line break (e.g. "gradua-\nted" -> "graduated")
  - Mid-paragraph line breaks (single \n) -> joined with a space
  - Paragraph breaks (blank line, or a line ending in real punctuation
    followed by a line that starts with a capital/indent) -> preserved
  - Leftover multiple blank lines collapsed to one

Usage as a script:
    python reflow_ocr.py input.txt output.txt
    python reflow_ocr.py input.txt          # prints to stdout
    cat input.txt | python reflow_ocr.py -  # read from stdin

Usage as a function:
    from reflow_ocr import reflow_text
    clean = reflow_text(raw_ocr_string)
"""

import re
import sys


# Words that legitimately end in a hyphen and should NOT be rejoined
# even though they're split at a line break (rare, but worth a hook).
# Add entries here if your source text has them, e.g. "self-", "co-".
HYPHEN_EXCEPTIONS = set()


def _dehyphenate(line_end_word: str, next_line_start_word: str) -> bool:
    """
    Decide whether a hyphenated line-break should be rejoined.
    Default heuristic: rejoin unless the word is a known exception.
    """
    return line_end_word.lower() not in HYPHEN_EXCEPTIONS


def reflow_text(text: str) -> str:
    """
    Reflow arbitrarily-wrapped OCR text into normal paragraphs.

    Paragraph breaks in the *input* are detected as blank lines
    (one or more empty lines). Everything else is treated as a
    soft wrap and joined into a single paragraph with spaces,
    rejoining hyphenated word-breaks as it goes.
    """
    # Normalize line endings
    text = text.replace('\r\n', '\n').replace('\r', '\n')

    # Split into paragraph blocks on one-or-more blank lines
    raw_paragraphs = re.split(r'\n\s*\n+', text)

    cleaned_paragraphs = []
    for block in raw_paragraphs:
        lines = [ln.strip() for ln in block.split('\n')]
        lines = [ln for ln in lines if ln != '']
        if not lines:
            continue

        paragraph = lines[0]
        for line in lines[1:]:
            # Case 1: previous chunk ends with a hyphen -> word was split
            hyphen_match = re.search(r'(\S*)-\Z', paragraph)
            if hyphen_match and re.match(r'^[a-z]', line):
                # word continues on next line, e.g. "gradua-" + "ted ..."
                next_word_match = re.match(r'^(\S+)', line)
                next_word = next_word_match.group(1) if next_word_match else ''
                broken_word = hyphen_match.group(1)
                if _dehyphenate(broken_word, next_word):
                    paragraph = paragraph[:-1] + line  # drop hyphen, glue together
                    continue
            # Case 2: normal soft wrap -> join with a space
            paragraph += ' ' + line

        # Collapse any accidental double spaces
        paragraph = re.sub(r' {2,}', ' ', paragraph).strip()
        cleaned_paragraphs.append(paragraph)

    return '\n\n'.join(cleaned_paragraphs)


def reflow_file(input_path: str, output_path: str = None) -> str:
    """Reflow a file's contents and optionally write the result to output_path."""
    if input_path == '-':
        raw = sys.stdin.read()
    else:
        with open(input_path, 'r', encoding='utf-8') as f:
            raw = f.read()

    cleaned = reflow_text(raw)

    if output_path:
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(cleaned)
    return cleaned


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    in_path = sys.argv[1]
    out_path = sys.argv[2] if len(sys.argv) > 2 else None

    result = reflow_file(in_path, out_path)

    if not out_path:
        print(result)