from __future__ import annotations
import argparse
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from extract_lib import extract_labeled_pcaps
DEFAULT_PCAP_ARCHIVE_DIR = Path('datasets/iscxtor2016/raw/pcap')
DEFAULT_DECOMPRESS_DIR = Path('datasets/iscxtor2016/raw/pcap_extracted')
DEFAULT_OUT_PACKETS = Path('datasets/iscxtor2016/processed/packets.npz')
DEFAULT_OUT_FLOWS = Path('datasets/iscxtor2016/processed/flows.parquet')
NONTOR_ARCHIVE = 'NonTor.tar.xz'
TOR_ARCHIVE = 'Tor.zip'
ACTIVITY_PATTERNS = (('mail', re.compile('mail|email|imap|pop_|smtp|thunderbird')), ('voip', re.compile('voip|voice|call|facebook_voice|hangouts_voice')), ('audio', re.compile('audio|spotify|skype_audio|hangout_audio|facebook_audio')), ('browsing', re.compile('browsing|browser|ssl_browsing|gate_ssl')), ('chat', re.compile('chat|aim|icq|skypechat')), ('file', re.compile('file[-_]?transfer|ftp|sftp|tftp')), ('p2p', re.compile('p2p|multispeed|multiple[sS]peed|bittor|utor')), ('video', re.compile('video|youtube|vimeo')))

def _infer_activity(pcap_name: str) -> str:
    lower = pcap_name.lower()
    for (act, pat) in ACTIVITY_PATTERNS:
        if pat.search(lower):
            return act
    return 'other'

def _decompress_archives(archive_dir: Path, out_dir: Path) -> None:
    nontor_arc = archive_dir / NONTOR_ARCHIVE
    tor_arc = archive_dir / TOR_ARCHIVE
    out_nontor = out_dir / 'NonTor'
    out_tor = out_dir / 'Tor'
    if not out_nontor.exists():
        out_nontor.parent.mkdir(parents=True, exist_ok=True)
        print(f'[decompress] {nontor_arc} → {out_dir}/  (tar xf)')
        t0 = time.time()
        subprocess.run(['tar', '-xf', str(nontor_arc), '-C', str(out_dir)], check=True)
        print(f'[decompress] NonTor done in {time.time() - t0:.1f}s')
    else:
        print(f'[decompress] {out_nontor} already exists — skipping NonTor unpack')
    if not out_tor.exists():
        print(f'[decompress] {tor_arc} → {out_dir}/  (unzip)')
        t0 = time.time()
        subprocess.run(['unzip', '-q', '-o', str(tor_arc), '-d', str(out_dir)], check=True)
        print(f'[decompress] Tor done in {time.time() - t0:.1f}s')
    else:
        print(f'[decompress] {out_tor} already exists — skipping Tor unpack')

def _find_pcap_files(decompressed_root: Path) -> list[tuple[Path, str, dict]]:
    triples: list[tuple[Path, str, dict]] = []
    for (sub, coarse) in (('NonTor', 'nontor'), ('Tor', 'tor')):
        sub_dir = decompressed_root / sub
        if not sub_dir.exists():
            print(f'[warn] {sub_dir} not found — skipping')
            continue
        pcaps = sorted(sub_dir.rglob('*.pcap')) + sorted(sub_dir.rglob('*.pcapng'))
        for p in pcaps:
            activity = _infer_activity(p.name)
            triples.append((p, coarse, {'activity': activity}))
    return triples

def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--archive-dir', type=Path, default=DEFAULT_PCAP_ARCHIVE_DIR)
    ap.add_argument('--decompressed-dir', type=Path, default=DEFAULT_DECOMPRESS_DIR)
    ap.add_argument('--out-packets', type=Path, default=DEFAULT_OUT_PACKETS)
    ap.add_argument('--out-flows', type=Path, default=DEFAULT_OUT_FLOWS)
    ap.add_argument('--out-store', type=Path, default=None, help='Optional sharded packet store output. When set, writes store_root/{metadata,manifest,flows,packets/*} instead of the monolithic packets.npz/flows.parquet pair.')
    ap.add_argument('--shard-size', type=int, default=100000, help='Rows per packet shard when --out-store is set.')
    ap.add_argument('--worker-flush-size', type=int, default=10000, help='Flows per temporary worker chunk when --out-store is set.')
    ap.add_argument('--spool-dir', type=Path, default=None, help='Optional temporary spool directory for worker chunks.')
    ap.add_argument('--T-full', type=int, default=256)
    ap.add_argument('--idle-timeout', type=float, default=120.0)
    ap.add_argument('--jobs', type=int, default=0)
    ap.add_argument('--max-packets-per-pcap', type=int, default=None)
    ap.add_argument('--decompress-only', action='store_true', help='Extract the archives then stop (for staged runs).')
    ap.add_argument('--skip-decompress', action='store_true', help='Assume decompressed-dir is already populated.')
    args = ap.parse_args()
    if not args.skip_decompress:
        _decompress_archives(args.archive_dir, args.decompressed_dir)
    if args.decompress_only:
        print('[decompress-only] exiting as requested.')
        return
    triples = _find_pcap_files(args.decompressed_dir)
    if not triples:
        raise RuntimeError(f'No pcap files found under {args.decompressed_dir}')
    print(f'\n[discover] found {len(triples)} pcap file(s)')
    by_coarse: dict[str, int] = {}
    by_act: dict[str, int] = {}
    for (_, lbl, extra) in triples:
        by_coarse[lbl] = by_coarse.get(lbl, 0) + 1
        by_act[extra['activity']] = by_act.get(extra['activity'], 0) + 1
    print(f'    by label: {by_coarse}')
    print(f'    by activity: {by_act}')
    print(f'\n[extract] writing to {args.out_packets} + {args.out_flows}')
    extract_labeled_pcaps(pcap_files_with_labels=triples, out_packets=args.out_packets, out_flows=args.out_flows, out_store=args.out_store, shard_size=args.shard_size, worker_flush_size=args.worker_flush_size, spool_dir=args.spool_dir, T_full=args.T_full, idle_timeout=args.idle_timeout, max_packets_per_pcap=args.max_packets_per_pcap, n_jobs=args.jobs, extra_column_names=('activity',))
if __name__ == '__main__':
    main()
