"""Batch generator: ~100 CSV files with varied column subsets for multi-upload testing.

Each file: 100 rows, 20-40% columns randomly dropped (IP/port identity columns retained).
Total: ~100 files × 100 rows = ~10000 rows.

Usage:
    runtime\\python\\python.exe scripts\\gen_multi.py [--files N] [--rows N]
"""
import sys
import os
import random
import csv
from pathlib import Path
from datetime import datetime, timedelta

# Add scripts/ to path so we can import gen_test_data
sys.path.insert(0, str(Path(__file__).parent))

# Reuse gen_test_data's data pools and generation logic
from gen_test_data import (
    COLUMNS, COVERAGE, generate_row, _rand_ip, _rand_hex, _maybe, _blank_or, _blank_or_plus,
    TLS_VERSIONS, TLS_WEIGHTS, CIPHER_HEX, CIPHER_NAMES, NAMED_CURVES_HEX, NAMED_CURVES_TEXT,
    SRC_IPS, DST_IPS, PORTS, COUNTRIES, ISPS, ORGS, CITIES, SERVICES, SRC_NODES, NODE_NAMES,
)


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Generate multi-file test CSV dataset')
    parser.add_argument('--files', type=int, default=100, help='Number of files to generate')
    parser.add_argument('--rows', type=int, default=100, help='Rows per file')
    parser.add_argument('--start', type=int, default=0, help='Start file index (for resume)')
    args = parser.parse_args()

    # Use a different seed from gen_test_data so IP pools differ
    random.seed(99)

    out_dir = Path('data/multi_upload')
    out_dir.mkdir(parents=True, exist_ok=True)

    # Identity columns that must never be dropped
    _keep_cols = {':ips', ':ipd', ':prs', ':prd', 'server-ip', 'client-ip'}
    _droppable = [c for c in COLUMNS if c not in _keep_cols]

    base_time = datetime(2026, 6, 1, 0, 0, 0)
    # Resume support: start from --start, estimate prior rows
    total_rows = args.start * args.rows if args.start > 0 else 0

    for fi in range(args.start, args.files):
        # Coverage-based column dropping: columns with lower coverage are more likely dropped.
        # Drop a column with probability (1 - coverage), capped to [0.10, 0.90] range.
        dropped = set()
        for col in _droppable:
            cov = COVERAGE.get(col, 1.0)
            drop_prob = 1.0 - cov
            # Clamp to ensure some columns always have a chance to be dropped/kept
            drop_prob = max(0.10, min(0.90, drop_prob))
            if random.random() < drop_prob:
                dropped.add(col)
        # Ensure at least 1 column is dropped and not ALL droppable columns
        if len(dropped) < 3:
            # Add a few extra random drops to ensure variety
            extra = random.sample([c for c in _droppable if c not in dropped],
                                  k=min(3 - len(dropped), len(_droppable) - len(dropped)))
            dropped.update(extra)
        if len(dropped) == len(_droppable):
            # Keep at least 3 droppable columns
            spared = set(random.sample(list(dropped), k=min(3, len(dropped))))
            dropped -= spared
        file_columns = [c for c in COLUMNS if c not in dropped]

        out_path = out_dir / f'{fi}.csv'
        n_rows = random.randint(max(50, args.rows - 10), args.rows + 10)

        with open(out_path, 'w', newline='', encoding='utf-8') as f:
            writer = csv.DictWriter(f, fieldnames=file_columns, extrasaction='ignore')
            writer.writeheader()
            for row_i in range(n_rows):
                row_time = base_time + timedelta(seconds=random.randint(0, 86400 * 30))
                row = generate_row(total_rows + row_i, row_time)
                filtered = {k: v for k, v in row.items() if k in file_columns}
                writer.writerow(filtered)

        total_rows += n_rows

        if (fi + 1) % 20 == 0:
            print(f'  [{fi + 1}/{args.files}] {total_rows} rows generated so far...')

    print(f'\nDone: {args.files} files, {total_rows} total rows → {out_dir}')
    print(f'Each file: ~{args.rows} rows, {len(_droppable)} droppable columns, coverage-based dropping per file')


if __name__ == '__main__':
    main()
