"""
Normalize extracted text before indexing.

Usage:
    python scripts/normalize_text.py <input_dir> <output_dir>

Reads JSON files produced by parse_wikipedia.py / parse_pdf.py,
cleans the text field, writes cleaned files to output_dir.
"""
import sys
import re
import json
import argparse
from pathlib import Path


def normalize(text: str) -> str:
    # Collapse unicode whitespace and control characters
    text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', text)

    # Normalize various dash/hyphen variants to a regular hyphen
    text = re.sub(r'[—–‒―]', '-', text)

    # Remove repeated punctuation (e.g. "!!!" → "!")
    text = re.sub(r'([!?.]){2,}', r'\1', text)

    # Collapse multiple spaces / tabs to a single space
    text = re.sub(r'[ \t]+', ' ', text)

    # Collapse more than two consecutive newlines
    text = re.sub(r'\n{3,}', '\n\n', text)

    # Strip leading/trailing whitespace per line
    lines = [line.strip() for line in text.splitlines()]
    text = '\n'.join(lines)

    return text.strip()


def process_directory(input_dir: str, output_dir: str) -> int:
    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    count = 0
    for src in sorted(Path(input_dir).glob('*.json')):
        try:
            with open(src, encoding='utf-8') as f:
                doc = json.load(f)
        except Exception as e:
            print(f'Skip {src.name}: {e}')
            continue

        if 'text' not in doc:
            continue

        doc['text'] = normalize(doc['text'])

        if len(doc['text'].split()) < 30:
            continue

        dest = out_path / src.name
        with open(dest, 'w', encoding='utf-8') as f:
            json.dump(doc, f, ensure_ascii=False, indent=2)

        count += 1
        if count % 1000 == 0:
            print(f'Normalized {count} documents')

    return count


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('input_dir')
    parser.add_argument('output_dir')
    args = parser.parse_args()

    count = process_directory(args.input_dir, args.output_dir)
    print(f'Done: {count} documents written to {args.output_dir}')


if __name__ == '__main__':
    main()