pdf_to_csv.py

#!/usr/bin/env python3
"""Extract text from a large PDF and write rows to a CSV with custom separator.

Strategy:
- Use PyMuPDF (fitz) to stream pages one-by-one (memory efficient)
- Split each non-empty line by a regex (default: two-or-more spaces) into columns
- Write rows incrementally to the CSV using the requested delimiter (default `;`)

Usage example:
  python3 pdf_to_csv.py input.pdf output.csv --sep ";" --split-pattern "\\s{2,}"
"""

from __future__ import annotations
import re
import argparse
import csv
import sys
import os
from typing import Iterable

try:
    import fitz  # PyMuPDF
except Exception:
    print("Missing dependency: PyMuPDF (fitz). Install via: pip install -r requirements.txt", file=sys.stderr)
    raise


def iter_lines_from_pdf(path: str, start: int | None = None, end: int | None = None) -> Iterable[str]:
    doc = fitz.open(path)
    page_count = doc.page_count
    s = 0 if start is None else max(0, start)
    e = page_count - 1 if end is None else min(page_count - 1, end)
    for pno in range(s, e + 1):
        page = doc.load_page(pno)
        text = page.get_text("text")
        for line in text.splitlines():
            yield line


def convert(pdf_path: str, out_csv: str, sep: str = ";", split_pattern: str = r"\s{2,}",
            start: int | None = None, end: int | None = None, encoding: str = "utf-8") -> None:
    splitter = re.compile(split_pattern)

    # Determine output path: if no directory provided, place into ./output_csv
    if os.path.isdir(out_csv) or out_csv.endswith(os.sep):
        base = os.path.splitext(os.path.basename(pdf_path))[0] + ".csv"
        out_path = os.path.join(out_csv, base)
    elif os.path.dirname(out_csv) == "":
        out_path = os.path.join("output_csv", out_csv)
    else:
        out_path = out_csv

    os.makedirs(os.path.dirname(out_path), exist_ok=True)

    # Read all lines lazily and group them into records based on contatore ID
    lines = []
    meter_re = re.compile(r"^[A-Z]{2,}[0-9]{3,}")
    for raw_line in iter_lines_from_pdf(pdf_path, start, end):
        line = raw_line.strip()
        if not line:
            continue
        lines.append(line)

    # Extract headers from the first lines (before any record starts)
    headers = []
    records: list[list[str]] = []
    current: list[str] = []
    
    for line in lines:
        if meter_re.match(line):
            # Found a record ID - all previous non-empty lines are headers
            if not headers:
                headers = [h.strip() for h in lines[:lines.index(line)]]
            if current:
                records.append(current)
            current = [line]
        else:
            if not current:
                # Still collecting headers before first record
                continue
            current.append(line)
    if current:
        records.append(current)
    
    # Fallback to hardcoded headers if extraction fails
    if not headers:
        headers = [
            "contatore",
            "contratto",
            "tecnologia",
            "anno",
            "data_rimozione",
            "guasto",
            "stato",
            "area",
            "polo",
            "ut_impresa",
            "indirizzo",
        ]

    with open(out_path, "w", newline="", encoding=encoding) as f:
        writer = csv.writer(f, delimiter=sep, quoting=csv.QUOTE_MINIMAL)
        writer.writerow(headers)
        for rec in records:
            # Initialize fields
            contatore = rec[0] if len(rec) > 0 else ""
            contratto = ""
            tecnologia = ""
            anno = ""
            data_rimozione = ""
            guasto = ""
            stato = ""
            area = ""
            polo = ""
            ut_impresa = ""
            indirizzo = ""

            idx = 1
            if len(rec) > idx:
                m = re.match(r"^(\d+)\s+(.*)$", rec[idx])
                if m:
                    contratto = m.group(1)
                    tecnologia = m.group(2)
                else:
                    parts = rec[idx].split(None, 1)
                    contratto = parts[0] if parts else ""
                    tecnologia = parts[1] if len(parts) > 1 else ""
                idx += 1

            if len(rec) > idx and re.match(r"^\d{4}$", rec[idx]):
                anno = rec[idx]
                idx += 1

            if len(rec) > idx and re.match(r"^\d{1,2}/\d{1,2}/\d{2,4}", rec[idx]):
                parts = rec[idx].split(None, 1)
                data_rimozione = parts[0]
                guasto = parts[1] if len(parts) > 1 else ""
                idx += 1

            if len(rec) > idx:
                stato = rec[idx]
                idx += 1

            if len(rec) > idx and re.match(r"^\(\d+\)\s*-", rec[idx]):
                area = rec[idx]
                idx += 1

            # Remaining lines: try to pick POLO, ut_impresa, indirizzo
            if len(rec) > idx:
                polo = rec[idx]
                idx += 1
            if len(rec) > idx:
                ut_impresa = rec[idx]
                idx += 1
            if len(rec) > idx:
                indirizzo = " ".join(rec[idx:])

            row = [
                contatore,
                contratto,
                tecnologia,
                anno,
                data_rimozione,
                guasto,
                stato,
                area,
                polo,
                ut_impresa,
                indirizzo,
            ]
            writer.writerow(row)


def get_default_pdf():
    """Get the first PDF from input_pdf folder, or None if not found."""
    input_dir = "input_pdf"
    if os.path.isdir(input_dir):
        pdfs = [f for f in os.listdir(input_dir) if f.lower().endswith('.pdf')]
        if pdfs:
            return os.path.join(input_dir, sorted(pdfs)[0])
    return None


def get_default_csv_name(pdf_path):
    """Generate output CSV filename based on PDF filename."""
    basename = os.path.splitext(os.path.basename(pdf_path))[0]
    return os.path.join("output_csv", f"{basename}.csv")


def build_parser() -> argparse.ArgumentParser:
    default_pdf = get_default_pdf()
    default_csv = get_default_csv_name(default_pdf) if default_pdf else None
    
    p = argparse.ArgumentParser(description="Convert PDF lines to CSV (delimiter configurable)")
    p.add_argument("pdf", nargs="?", default=default_pdf, help=f"Input PDF file path (default: {default_pdf})")
    p.add_argument("csv", nargs="?", default=default_csv, help=f"Output CSV file path (default: {default_csv})")
    p.add_argument("--sep", default=";", help="Output CSV delimiter (default: ';')")
    p.add_argument("--split-pattern", default=r"\s{2,}", help="Regex pattern to split a line into columns (default: two or more spaces)")
    p.add_argument("--start-page", type=int, default=None, help="0-based start page (inclusive)")
    p.add_argument("--end-page", type=int, default=None, help="0-based end page (inclusive)")
    p.add_argument("--encoding", default="utf-8", help="Output file encoding (default utf-8)")
    return p


def generate_config_json(csv_path: str, sep: str = ";"):
    """Generate config.json from CSV headers."""
    try:
        import json
        headers = []
        with open(csv_path, 'r', encoding='utf-8') as f:
            reader = csv.reader(f, delimiter=sep)
            headers = next(reader, [])
        
        if headers:
            config = {
                "input_csv": csv_path,
                "output_dir": "extracted_csv",
                "filter_column": headers[0],
                "output_filename_column": headers[0],
                "values_to_filter": [],
                "csv_separator": sep,
                "available_columns": headers,
                "description": f"Configuration file for extract_contracts.py. Available columns: {', '.join(headers)}"
            }
            
            config_path = "config.json"
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
            print(f"✓ Config generated: {config_path}")
            print(f"  Available columns: {', '.join(headers)}")
    except Exception as e:
        print(f"Warning: Could not generate config.json: {e}", file=sys.stderr)


def main(argv: list[str] | None = None) -> int:
    p = build_parser()
    args = p.parse_args(argv)
    try:
        convert(args.pdf, args.csv, sep=args.sep, split_pattern=args.split_pattern,
                start=args.start_page, end=args.end_page, encoding=args.encoding)
        # Generate config.json after successful conversion
        generate_config_json(args.csv, args.sep)
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())