#!/usr/bin/env python3
"""Split a concatenated all_onion_texts.txt into per-source files.
Writes files into prepared/split_files/ using sanitized header names.
Usage: python scripts/split_all_onion_texts.py --in-file prepared/all_onion_texts.txt --out-dir prepared/split_files --max 200
"""
import argparse
import os
from pathlib import Path

def sanitize_name(s: str) -> str:
    # keep ascii and replace path separators with underscores
    s = s.strip()
    s = s.replace(':', '')
    s = s.replace('\\', '_').replace('/', '_')
    s = s.replace(' ', '_')
    # remove suspicious characters
    for ch in ['"', "'", '<', '>', '|', '?', '*', ':']:
        s = s.replace(ch, '')
    return s[:200]


def split_file(in_file: Path, out_dir: Path, max_files: int = None):
    out_dir.mkdir(parents=True, exist_ok=True)
    current_f = None
    current_name = None
    file_count = 0
    header_prefix = '--- FILE:'
    with in_file.open('r', errors='ignore') as fh:
        for line in fh:
            if line.startswith(header_prefix):
                # start a new file
                # header format: --- FILE: <path> ---
                try:
                    header = line.strip()
                    # extract between prefix and closing ---
                    if '---' in header[len(header_prefix):]:
                        # remove prefix
                        rest = header[len(header_prefix):].strip()
                        # remove trailing ---
                        if rest.endswith('---'):
                            rest = rest[:-3].strip()
                        name = sanitize_name(rest)
                    else:
                        name = sanitize_name(header[len(header_prefix):])
                except Exception:
                    name = f"split_{file_count}"
                if current_f:
                    current_f.close()
                if max_files is not None and file_count >= max_files:
                    break
                outfile = out_dir / f"{file_count:05d}_{name}.txt"
                current_f = outfile.open('w', errors='ignore')
                current_name = name
                file_count += 1
            else:
                if current_f:
                    current_f.write(line)
    if current_f:
        current_f.close()
    return file_count


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--in-file', required=True)
    parser.add_argument('--out-dir', required=True)
    parser.add_argument('--max', type=int, default=200, help='maximum number of files to write (default 200). Use 0 or omit to split all')
    args = parser.parse_args()
    in_file = Path(args.in_file)
    out_dir = Path(args.out_dir)
    mx = args.max if args.max and args.max > 0 else None
    print('Splitting', in_file, '->', out_dir, 'max=', mx)
    n = split_file(in_file, out_dir, max_files=mx)
    print('Wrote', n, 'files')